From 22e344be15c6e505fc33cc7c74ceed4c654cf0be Mon Sep 17 00:00:00 2001 From: Joe Nudell Date: Fri, 10 Jul 2026 12:35:05 -0400 Subject: [PATCH 1/3] feat: track usage and estimate cost for calls --- README.md | 32 +++ bc2/core/analyze/azuredi.py | 19 +- bc2/core/analyze/test_azuredi.py | 36 +++ bc2/core/common/azure_pricing.py | 362 ++++++++++++++++++++++++++ bc2/core/common/openai.py | 46 ++++ bc2/core/common/pipe.py | 4 +- bc2/core/common/test_azure_pricing.py | 156 +++++++++++ bc2/core/common/test_openai.py | 47 +++- bc2/core/common/test_usage.py | 62 +++++ bc2/core/common/usage.py | 126 +++++++++ bc2/core/control/chunk.py | 4 +- bc2/core/pipeline.py | 8 +- bc2/core/test_pipeline.py | 14 + bc2/lib/embedding/openai.py | 24 +- bc2/lib/embedding/test_openai.py | 32 +++ 15 files changed, 965 insertions(+), 7 deletions(-) create mode 100644 bc2/core/analyze/test_azuredi.py create mode 100644 bc2/core/common/azure_pricing.py create mode 100644 bc2/core/common/test_azure_pricing.py create mode 100644 bc2/core/common/test_usage.py create mode 100644 bc2/core/common/usage.py diff --git a/README.md b/README.md index cf618634..65d97cf5 100644 --- a/README.md +++ b/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. diff --git a/bc2/core/analyze/azuredi.py b/bc2/core/analyze/azuredi.py index 605b1887..ba3cb29b 100644 --- a/bc2/core/analyze/azuredi.py +++ b/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]() diff --git a/bc2/core/analyze/test_azuredi.py b/bc2/core/analyze/test_azuredi.py new file mode 100644 index 00000000..ef9f805b --- /dev/null +++ b/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} diff --git a/bc2/core/common/azure_pricing.py b/bc2/core/common/azure_pricing.py new file mode 100644 index 00000000..b80963b3 --- /dev/null +++ b/bc2/core/common/azure_pricing.py @@ -0,0 +1,362 @@ +from __future__ import annotations + +import json +import re +import threading +import time +import urllib.parse +import urllib.request +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any + +AZURE_RETAIL_PRICES_URL = "https://prices.azure.com/api/retail/prices" +DEFAULT_CACHE_TTL_SECONDS = 24 * 60 * 60 + + +class AzurePricingUnavailable(Exception): + """Raised when an Azure retail price cannot be determined unambiguously.""" + + +@dataclass +class _CacheEntry: + expires_at: float + fetched_at: str + items: list[dict[str, Any]] + + +class AzureRetailPricing: + """Fetch and cache public Azure retail prices.""" + + def __init__(self, cache_ttl_seconds: int = DEFAULT_CACHE_TTL_SECONDS): + self.cache_ttl_seconds = cache_ttl_seconds + self._cache: dict[str, _CacheEntry] = {} + self._lock = threading.Lock() + + def estimate( + self, call: dict[str, Any], runtime_config: dict[str, Any] + ) -> dict[str, Any]: + """Estimate the cost of one Azure service call.""" + region = runtime_config.get("azure_region") + if not region: + raise AzurePricingUnavailable( + "azure_region is required to look up Azure retail pricing" + ) + + service = call.get("service") + if service == "responses": + result = self._estimate_openai_tokens(call, runtime_config, region) + elif service == "embeddings": + result = self._estimate_openai_tokens(call, runtime_config, region) + elif service == "document_intelligence": + result = self._estimate_document_intelligence(call, region) + result["sku"] = "S0" + else: + raise AzurePricingUnavailable(f"unsupported Azure service: {service}") + + result["region"] = region + if service in {"responses", "embeddings"}: + result["deployment_type"] = runtime_config.get( + "azure_deployment_type", "global" + ) + result["context_tier"] = runtime_config.get("azure_context_tier", "short") + return result + + def _estimate_openai_tokens( + self, + call: dict[str, Any], + runtime_config: dict[str, Any], + region: str, + ) -> dict[str, Any]: + model = call.get("model") + if not model: + raise AzurePricingUnavailable("OpenAI model was not reported") + + is_embedding = call.get("service") == "embeddings" + search_term = "embedding" if is_embedding else _model_search_term(model) + odata_filter = ( + "contains(productName, 'OpenAI') " + f"and armRegionName eq '{_odata_escape(region)}' " + f"and contains(skuName, '{_odata_escape(search_term)}')" + ) + items, fetched_at = self._get_prices(odata_filter) + deployment_type = runtime_config.get("azure_deployment_type", "global") + context_tier = runtime_config.get("azure_context_tier", "short") + usage = call.get("usage", {}) + + components: list[dict[str, Any]] = [] + input_tokens = int(usage.get("input_tokens") or 0) + cached_tokens = int(usage.get("cached_input_tokens") or 0) + regular_input_tokens = max(0, input_tokens - cached_tokens) + output_tokens = int(usage.get("output_tokens") or 0) + + if is_embedding and input_tokens: + meter = _select_embedding_meter( + items, + model=model, + deployment_type=deployment_type, + ) + components.append(_price_tokens(input_tokens, meter, "input")) + return _cost_result(components, fetched_at) + + if regular_input_tokens: + meter = _select_openai_meter( + items, + model=model, + kind="input", + deployment_type=deployment_type, + context_tier=context_tier, + ) + components.append(_price_tokens(regular_input_tokens, meter, "input")) + if cached_tokens: + meter = _select_openai_meter( + items, + model=model, + kind="cached_input", + deployment_type=deployment_type, + context_tier=context_tier, + ) + components.append(_price_tokens(cached_tokens, meter, "cached_input")) + if output_tokens: + meter = _select_openai_meter( + items, + model=model, + kind="output", + deployment_type=deployment_type, + context_tier=context_tier, + ) + components.append(_price_tokens(output_tokens, meter, "output")) + + if not components: + raise AzurePricingUnavailable("no billable token usage was reported") + return _cost_result(components, fetched_at) + + def _estimate_document_intelligence( + self, call: dict[str, Any], region: str + ) -> dict[str, Any]: + model = str(call.get("model") or "") + features = call.get("features") or [] + if model != "prebuilt-read": + raise AzurePricingUnavailable( + f"unsupported Document Intelligence model: {model or ''}" + ) + if features: + raise AzurePricingUnavailable( + "Document Intelligence feature add-on pricing is not supported" + ) + + pages = int((call.get("usage") or {}).get("pages") or 0) + if pages <= 0: + raise AzurePricingUnavailable("no analyzed pages were reported") + + odata_filter = ( + "productName eq 'Azure Document Intelligence' " + f"and armRegionName eq '{_odata_escape(region)}' " + "and skuName eq 'S0'" + ) + items, fetched_at = self._get_prices(odata_filter) + matches = [ + item + for item in items + if item.get("meterName") == "S0 Read Pages" + and item.get("type") == "Consumption" + and float(item.get("tierMinimumUnits") or 0) == 0 + ] + meter = _select_unique_price(matches, "S0 Read Pages") + component = _price_quantity(pages, meter, "pages") + return _cost_result([component], fetched_at) + + def _get_prices(self, odata_filter: str) -> tuple[list[dict[str, Any]], str]: + now = time.monotonic() + with self._lock: + cached = self._cache.get(odata_filter) + if cached and cached.expires_at > now: + return cached.items, cached.fetched_at + + params = urllib.parse.urlencode( + { + "api-version": "2023-01-01-preview", + "$filter": odata_filter, + "currencyCode": "USD", + } + ) + url: str | None = f"{AZURE_RETAIL_PRICES_URL}?{params}" + items: list[dict[str, Any]] = [] + while url: + with urllib.request.urlopen(url, timeout=10) as response: + payload = json.load(response) + items.extend(payload.get("Items") or []) + url = payload.get("NextPageLink") + + fetched_at = datetime.now(timezone.utc).isoformat() + entry = _CacheEntry( + expires_at=now + self.cache_ttl_seconds, + fetched_at=fetched_at, + items=items, + ) + with self._lock: + self._cache[odata_filter] = entry + return items, fetched_at + + +def _model_search_term(model: str) -> str: + normalized = re.sub(r"-\d{4}-\d{2}-\d{2}$", "", model.lower()) + if normalized.startswith("gpt-"): + return normalized.removeprefix("gpt-") + return normalized + + +def _odata_escape(value: str) -> str: + return value.replace("'", "''") + + +def _select_openai_meter( + items: list[dict[str, Any]], + *, + model: str, + kind: str, + deployment_type: str, + context_tier: str, +) -> dict[str, Any]: + matches = [] + for item in items: + if item.get("type") != "Consumption": + continue + name = f"{item.get('skuName', '')} {item.get('meterName', '')}".lower() + if "batch" in name or re.search(r"\bpp\b", name): + continue + if not _matches_model_variant(name, model): + continue + if not _matches_deployment_type(name, deployment_type): + continue + if not _matches_context_tier(name, context_tier): + continue + + is_cached = bool(re.search(r"\b(cached|cchd|cd)\b", name)) + is_input = bool(re.search(r"\b(input|inpt|inp)\b", name)) + is_output = bool(re.search(r"\b(output|outp|opt)\b", name)) + if kind == "input" and is_input and not is_cached: + matches.append(item) + elif kind == "cached_input" and is_input and is_cached: + matches.append(item) + elif kind == "output" and is_output and not is_cached: + matches.append(item) + + return _select_unique_price(matches, f"OpenAI {kind}") + + +def _matches_model_variant(name: str, model: str) -> bool: + normalized_model = re.sub(r"-\d{4}-\d{2}-\d{2}$", "", model.lower()) + qualifiers = ("mini", "nano", "ft", "dev", "audio", "realtime", "codex", "pro") + for qualifier in qualifiers: + model_has_qualifier = bool( + re.search(rf"(^|[-_. ]){qualifier}($|[-_. ])", normalized_model) + ) + name_has_qualifier = bool(re.search(rf"(^|[-_. ]){qualifier}($|[-_. ])", name)) + if model_has_qualifier != name_has_qualifier: + return False + return "grader" not in name + + +def _select_embedding_meter( + items: list[dict[str, Any]], + *, + model: str, + deployment_type: str, +) -> dict[str, Any]: + target = re.sub(r"[^a-z0-9]", "", model.lower()) + matches = [] + for item in items: + if item.get("type") != "Consumption": + continue + name = f"{item.get('skuName', '')} {item.get('meterName', '')}".lower() + compact_name = re.sub(r"[^a-z0-9]", "", name) + if target not in compact_name: + continue + if "grader" in name or not _matches_deployment_type(name, deployment_type): + continue + matches.append(item) + return _select_unique_price(matches, "OpenAI embedding input") + + +def _matches_deployment_type(name: str, deployment_type: str) -> bool: + is_global = bool(re.search(r"\b(global|glbl|gl)\b", name)) + is_data_zone = bool(re.search(r"\b(data zone|dzn|dz)\b", name)) + is_regional = bool(re.search(r"\b(regional|regnl)\b", name)) + if deployment_type == "global": + return is_global + if deployment_type == "data_zone": + return is_data_zone + if deployment_type == "regional": + return is_regional or (not is_global and not is_data_zone) + raise AzurePricingUnavailable( + "azure_deployment_type must be global, data_zone, or regional" + ) + + +def _matches_context_tier(name: str, context_tier: str) -> bool: + is_short = "shortco" in name + is_long = "longco" in name + if not is_short and not is_long: + return True + if context_tier == "short": + return is_short + if context_tier == "long": + return is_long + raise AzurePricingUnavailable("azure_context_tier must be short or long") + + +def _select_unique_price( + matches: list[dict[str, Any]], description: str +) -> dict[str, Any]: + prices = { + (float(item["retailPrice"]), str(item["unitOfMeasure"])) for item in matches + } + if not prices: + raise AzurePricingUnavailable(f"no retail meter found for {description}") + if len(prices) > 1: + raise AzurePricingUnavailable( + f"multiple retail meters matched {description}; pricing is ambiguous" + ) + return matches[0] + + +def _price_tokens(tokens: int, meter: dict[str, Any], category: str) -> dict[str, Any]: + return _price_quantity(tokens, meter, category) + + +def _price_quantity( + quantity: int, meter: dict[str, Any], category: str +) -> dict[str, Any]: + unit = str(meter["unitOfMeasure"]) + divisor = _unit_divisor(unit) + unit_price = float(meter["retailPrice"]) + return { + "category": category, + "quantity": quantity, + "meter_name": meter["meterName"], + "unit_of_measure": unit, + "unit_price": unit_price, + "estimated_cost": quantity / divisor * unit_price, + } + + +def _unit_divisor(unit: str) -> int: + normalized = unit.upper() + if normalized == "1K": + return 1_000 + if normalized == "1M": + return 1_000_000 + if normalized == "1": + return 1 + raise AzurePricingUnavailable(f"unsupported Azure pricing unit: {unit}") + + +def _cost_result(components: list[dict[str, Any]], fetched_at: str) -> dict[str, Any]: + return { + "estimated_cost": sum(c["estimated_cost"] for c in components), + "currency": "USD", + "source": "azure-retail-prices", + "pricing_fetched_at": fetched_at, + "components": components, + } diff --git a/bc2/core/common/openai.py b/bc2/core/common/openai.py index bf27bacc..679fbf29 100644 --- a/bc2/core/common/openai.py +++ b/bc2/core/common/openai.py @@ -24,6 +24,7 @@ from .image import ImageUrl from .openai_metadata import ModelNotFound, get_chat_model_meta from .template import TemplateEngine, get_formatter +from .usage import record_usage logger = logging.getLogger(__name__) @@ -527,6 +528,8 @@ def invoke( else: response = client.responses.create(**call_params) + _record_response_usage(client, self, response) + # Interpret response stop_reason = response.incomplete_details and response.incomplete_details.reason truncated = False @@ -565,5 +568,48 @@ def invoke( ) +def _record_response_usage( + client: OpenAI, config: OpenAIChatConfig, response: Any +) -> None: + usage = getattr(response, "usage", None) + if usage is None: + return + + token_usage = { + "input_tokens": getattr(usage, "input_tokens", None), + "output_tokens": getattr(usage, "output_tokens", None), + "total_tokens": getattr(usage, "total_tokens", None), + } + input_details = getattr(usage, "input_tokens_details", None) + output_details = getattr(usage, "output_tokens_details", None) + token_usage["cached_input_tokens"] = getattr(input_details, "cached_tokens", None) + token_usage["reasoning_output_tokens"] = getattr( + output_details, "reasoning_tokens", None + ) + token_usage = {k: v for k, v in token_usage.items() if isinstance(v, int)} + + provider = _openai_provider(client) + reported_model = getattr(response, "model", None) + record_usage( + { + "provider": provider, + "service": "responses", + "model": config.openai_model + or (reported_model if isinstance(reported_model, str) else config.model), + "deployment": config.model if provider == "azure" else None, + "response_id": getattr(response, "id", None), + "status": getattr(response, "status", None), + "usage": token_usage, + } + ) + + +def _openai_provider(client: OpenAI | AsyncOpenAI) -> str: + base_url = str(getattr(client, "base_url", "")) + if "/openai/" in base_url or ".openai.azure.com" in base_url: + return "azure" + return "openai" + + class OpenAIConfig(BaseModel): client: OpenAIClientConfig diff --git a/bc2/core/common/pipe.py b/bc2/core/common/pipe.py index aebb2338..d8263f97 100644 --- a/bc2/core/common/pipe.py +++ b/bc2/core/common/pipe.py @@ -4,6 +4,7 @@ from .all import AnyConfig from .context import Context from .type_util import inspect_all_params, inspect_required_params, inspect_return_type +from .usage import usage_operation logger = logging.getLogger(__name__) @@ -153,7 +154,8 @@ def run_pipe( # NOTE(jnu): mypy can't validate the kwarg types, but we've effectively # done this at runtime anyway so just hush the error. try: - output = config.driver(*args, **kwargs) # type: ignore[arg-type] + with usage_operation(config.engine): + output = config.driver(*args, **kwargs) # type: ignore[arg-type] except Exception as e: logger.error(f"Error in step {i} ({config.engine}): {e}") ctx.errors.append(e) diff --git a/bc2/core/common/test_azure_pricing.py b/bc2/core/common/test_azure_pricing.py new file mode 100644 index 00000000..44f073ea --- /dev/null +++ b/bc2/core/common/test_azure_pricing.py @@ -0,0 +1,156 @@ +import io +import json + +import pytest + +from .azure_pricing import AzurePricingUnavailable, AzureRetailPricing + + +def _meter( + name: str, + price: float, + unit: str = "1M", + tier: float = 0, +) -> dict: + return { + "skuName": name, + "meterName": f"{name} Tokens", + "retailPrice": price, + "unitOfMeasure": unit, + "tierMinimumUnits": tier, + "type": "Consumption", + } + + +def test_estimate_openai_response_cost(monkeypatch): + pricing = AzureRetailPricing() + meters = [ + _meter("5.5 ShortCo inp Gl", 5.0), + _meter("5.5 ShortCo cd inp Gl", 0.5), + _meter("5.5 ShortCo opt Gl", 30.0), + ] + monkeypatch.setattr( + pricing, "_get_prices", lambda _: (meters, "2026-07-10T00:00:00+00:00") + ) + + estimate = pricing.estimate( + { + "service": "responses", + "model": "gpt-5.5", + "usage": { + "input_tokens": 1_000_000, + "cached_input_tokens": 100_000, + "output_tokens": 10_000, + }, + }, + { + "azure_region": "eastus", + "azure_deployment_type": "global", + "azure_context_tier": "short", + }, + ) + + assert estimate["estimated_cost"] == pytest.approx(4.85) + assert estimate["currency"] == "USD" + assert len(estimate["components"]) == 3 + + +def test_estimate_document_intelligence_read_cost(monkeypatch): + pricing = AzureRetailPricing() + meter = { + "skuName": "S0", + "meterName": "S0 Read Pages", + "retailPrice": 1.5, + "unitOfMeasure": "1K", + "tierMinimumUnits": 0, + "type": "Consumption", + } + monkeypatch.setattr( + pricing, "_get_prices", lambda _: ([meter], "2026-07-10T00:00:00+00:00") + ) + + estimate = pricing.estimate( + { + "service": "document_intelligence", + "model": "prebuilt-read", + "features": [], + "usage": {"pages": 25}, + }, + {"azure_region": "eastus"}, + ) + + assert estimate["estimated_cost"] == pytest.approx(0.0375) + + +def test_estimate_embedding_cost(monkeypatch): + pricing = AzureRetailPricing() + meter = _meter("text-embedding-3-large-glbl", 0.00013, "1K") + monkeypatch.setattr( + pricing, "_get_prices", lambda _: ([meter], "2026-07-10T00:00:00+00:00") + ) + + estimate = pricing.estimate( + { + "service": "embeddings", + "model": "text-embedding-3-large", + "usage": {"input_tokens": 2_000}, + }, + {"azure_region": "eastus", "azure_deployment_type": "global"}, + ) + + assert estimate["estimated_cost"] == pytest.approx(0.00026) + + +def test_missing_region_fails_without_fetching_prices(): + pricing = AzureRetailPricing() + + with pytest.raises(AzurePricingUnavailable, match="azure_region"): + pricing.estimate( + { + "service": "responses", + "model": "gpt-4.1", + "usage": {"input_tokens": 10}, + }, + {}, + ) + + +def test_ambiguous_meter_fails_gracefully(monkeypatch): + pricing = AzureRetailPricing() + meters = [ + _meter("4.1 Inp glbl", 2.0, "1K"), + _meter("4.1 Inp global", 3.0, "1K"), + ] + monkeypatch.setattr( + pricing, "_get_prices", lambda _: (meters, "2026-07-10T00:00:00+00:00") + ) + + with pytest.raises(AzurePricingUnavailable, match="ambiguous"): + pricing.estimate( + { + "service": "responses", + "model": "gpt-4.1", + "usage": {"input_tokens": 10}, + }, + {"azure_region": "eastus"}, + ) + + +def test_retail_prices_are_cached(monkeypatch): + pricing = AzureRetailPricing() + requests = 0 + + def fake_urlopen(url, timeout): + nonlocal requests + requests += 1 + return io.BytesIO( + json.dumps({"Items": [{"meterName": "test"}]}).encode("utf-8") + ) + + monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) + + first = pricing._get_prices("productName eq 'test'") + second = pricing._get_prices("productName eq 'test'") + + assert first == second + assert requests == 1 diff --git a/bc2/core/common/test_openai.py b/bc2/core/common/test_openai.py index 9d1d2cd1..ffb71b33 100644 --- a/bc2/core/common/test_openai.py +++ b/bc2/core/common/test_openai.py @@ -15,6 +15,7 @@ OpenAIChatTurn, OpenAIClientConfig, ) +from .usage import create_usage_tracker, usage_operation, usage_tracking def test_fix_azure_endpoint(): @@ -217,7 +218,21 @@ def _mock_response( response = MagicMock() response.status = status response.output_text = output_text - response.usage = type("Usage", (), {"output_tokens": output_tokens})() + response.usage = type( + "Usage", + (), + { + "input_tokens": 20, + "output_tokens": output_tokens, + "total_tokens": 20 + output_tokens, + "input_tokens_details": type("InputDetails", (), {"cached_tokens": 5})(), + "output_tokens_details": type( + "OutputDetails", (), {"reasoning_tokens": 3} + )(), + }, + )() + response.id = "resp_test" + response.model = "gpt-4.1-2025-04-14" if incomplete_reason is None: response.incomplete_details = None else: @@ -298,6 +313,36 @@ def test_invoke_completed_response_is_not_truncated(): assert result.is_truncated is False +def test_invoke_records_response_usage(): + cfg = _build_chat_config() + client = MagicMock() + client.base_url = "https://example.openai.azure.com/openai/v1/" + client.responses.create.return_value = _mock_response( + status="completed", + output_text="full answer", + output_tokens=42, + ) + created = create_usage_tracker({"report_usage": True}) + assert created is not None + report, tracker = created + + with usage_tracking(tracker), usage_operation("parse:openai"): + cfg.invoke(client, "go") + + call = report["calls"][0] + assert call["provider"] == "azure" + assert call["service"] == "responses" + assert call["operation"] == "parse:openai" + assert call["response_id"] == "resp_test" + assert call["usage"] == { + "input_tokens": 20, + "output_tokens": 42, + "total_tokens": 62, + "cached_input_tokens": 5, + "reasoning_output_tokens": 3, + } + + def test_chat_output_is_truncated_inferred_from_token_match(): """When the API doesn't surface a stop reason but we used every token, we still report the output as truncated.""" diff --git a/bc2/core/common/test_usage.py b/bc2/core/common/test_usage.py new file mode 100644 index 00000000..c2b18ef4 --- /dev/null +++ b/bc2/core/common/test_usage.py @@ -0,0 +1,62 @@ +from .usage import ( + create_usage_tracker, + record_usage, + usage_operation, + usage_tracking, +) + + +def test_usage_reporting_is_disabled_by_default(): + assert create_usage_tracker({}) is None + + +def test_estimate_cost_enables_usage_reporting(): + created = create_usage_tracker({"estimate_cost": True}) + + assert created is not None + + +def test_usage_call_and_totals_are_recorded(): + created = create_usage_tracker({"report_usage": True}) + assert created is not None + report, tracker = created + + with usage_tracking(tracker), usage_operation("redact:openai"): + record_usage( + { + "provider": "azure", + "service": "responses", + "model": "gpt-4.1", + "usage": { + "input_tokens": 100, + "output_tokens": 20, + "total_tokens": 120, + }, + } + ) + + assert report["calls"][0]["operation"] == "redact:openai" + assert report["totals"]["calls"] == 1 + assert report["totals"]["input_tokens"] == 100 + assert report["totals"]["output_tokens"] == 20 + + +def test_pricing_failure_is_added_to_call_instead_of_raised(): + created = create_usage_tracker({"estimate_cost": True}) + assert created is not None + report, tracker = created + + with usage_tracking(tracker): + record_usage( + { + "provider": "azure", + "service": "responses", + "model": "gpt-4.1", + "usage": {"input_tokens": 100}, + } + ) + + estimate = report["calls"][0]["cost_estimate"] + assert estimate["estimated_cost"] is None + assert "azure_region" in estimate["error"] + assert report["totals"]["unpriced_calls"] == 1 diff --git a/bc2/core/common/usage.py b/bc2/core/common/usage.py new file mode 100644 index 00000000..b4a11664 --- /dev/null +++ b/bc2/core/common/usage.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import logging +from contextlib import contextmanager +from contextvars import ContextVar +from threading import Lock +from typing import Any, Iterator + +from .azure_pricing import AzurePricingUnavailable, AzureRetailPricing + +logger = logging.getLogger(__name__) + +_current_tracker: ContextVar[UsageTracker | None] = ContextVar( + "bc2_usage_tracker", default=None +) +_current_operation: ContextVar[str | None] = ContextVar( + "bc2_usage_operation", default=None +) +_azure_pricing = AzureRetailPricing() + + +class UsageTracker: + """Collect per-call service usage for a single pipeline run.""" + + def __init__( + self, + report: dict[str, Any], + runtime_config: dict[str, Any], + ): + self.report = report + self.runtime_config = runtime_config + self.estimate_cost = bool(runtime_config.get("estimate_cost", False)) + self._lock = Lock() + + def record(self, call: dict[str, Any]) -> None: + """Record one service call without disrupting the pipeline.""" + call["operation"] = call.get("operation") or _current_operation.get() + if self.estimate_cost: + call["cost_estimate"] = self._estimate(call) + + with self._lock: + self.report["calls"].append(call) + totals = self.report["totals"] + totals["calls"] += 1 + for name, value in (call.get("usage") or {}).items(): + if isinstance(value, int): + totals[name] = totals.get(name, 0) + value + + estimate = call.get("cost_estimate") or {} + cost = estimate.get("estimated_cost") + if isinstance(cost, int | float): + totals["estimated_cost"] += cost + totals["estimated_calls"] += 1 + elif self.estimate_cost: + totals["unpriced_calls"] += 1 + + def _estimate(self, call: dict[str, Any]) -> dict[str, Any]: + if call.get("provider") != "azure": + return { + "estimated_cost": None, + "currency": "USD", + "error": "cost estimation is currently supported only for Azure", + } + try: + return _azure_pricing.estimate(call, self.runtime_config) + except AzurePricingUnavailable as exc: + return { + "estimated_cost": None, + "currency": "USD", + "error": str(exc), + } + except Exception as exc: + logger.warning("Unable to fetch Azure pricing: %s", exc) + return { + "estimated_cost": None, + "currency": "USD", + "error": f"Azure pricing lookup failed: {exc}", + } + + +def create_usage_tracker( + runtime_config: dict[str, Any], +) -> tuple[dict[str, Any], UsageTracker] | None: + """Create a usage report when runtime reporting is enabled.""" + if not ( + runtime_config.get("report_usage", False) + or runtime_config.get("estimate_cost", False) + ): + return None + report: dict[str, Any] = { + "calls": [], + "totals": { + "calls": 0, + "estimated_cost": 0.0, + "estimated_calls": 0, + "unpriced_calls": 0, + }, + } + return report, UsageTracker(report, runtime_config) + + +@contextmanager +def usage_tracking(tracker: UsageTracker | None) -> Iterator[None]: + """Make a pipeline's usage tracker available to service wrappers.""" + token = _current_tracker.set(tracker) + try: + yield + finally: + _current_tracker.reset(token) + + +@contextmanager +def usage_operation(operation: str) -> Iterator[None]: + """Attribute service calls to a pipeline operation.""" + token = _current_operation.set(operation) + try: + yield + finally: + _current_operation.reset(token) + + +def record_usage(call: dict[str, Any]) -> None: + """Record usage on the active pipeline tracker, if any.""" + tracker = _current_tracker.get() + if tracker is not None: + tracker.record(call) diff --git a/bc2/core/control/chunk.py b/bc2/core/control/chunk.py index 2e1e3e16..d5e3b1e3 100644 --- a/bc2/core/control/chunk.py +++ b/bc2/core/control/chunk.py @@ -15,6 +15,7 @@ inspect_required_params, inspect_return_type, ) +from ..common.usage import usage_operation from ..parse import ParseConfig from ..redact import RedactConfig from .compose import ComposeConfig @@ -139,7 +140,8 @@ def __call__( # Run the processor on the current chunk filtered_kwargs = get_bindable_parameters(f, runtime_config or {}) - new_output = cast(T, f(remainder, context, **filtered_kwargs)) + with usage_operation(self.config.processor.engine): + new_output = cast(T, f(remainder, context, **filtered_kwargs)) sep = " " if not context.debug else "\n\n\n|CHUNK BOUNDARY|\n\n\n" output = self._merge_output(output, new_output, separator=sep) diff --git a/bc2/core/pipeline.py b/bc2/core/pipeline.py index b4401129..370c862b 100644 --- a/bc2/core/pipeline.py +++ b/bc2/core/pipeline.py @@ -4,6 +4,7 @@ from .common.context import Context from .common.runtime import RuntimeConfig +from .common.usage import create_usage_tracker, usage_tracking logger = logging.getLogger(__name__) @@ -51,6 +52,10 @@ def run( ctx = Context() ctx.debug = runtime_config.get("debug", False) ctx.errors = list[Exception]() + usage = create_usage_tracker(runtime_config) + tracker = None + if usage is not None: + ctx.usage, tracker = usage if ctx.debug: # Set the global logger to info mode. logging.getLogger().setLevel(logging.INFO) @@ -61,7 +66,8 @@ def run( logger.debug("Debug mode enabled.") runtime_config["context"] = ctx - output = self.pipeline(None, ctx, runtime_config) + with usage_tracking(tracker): + output = self.pipeline(None, ctx, runtime_config) # The final pipe value is validated as None via type-checking. # It's not an error if the final pipe is not None, but we should log it. diff --git a/bc2/core/test_pipeline.py b/bc2/core/test_pipeline.py index dd273e7e..dc30c8e6 100644 --- a/bc2/core/test_pipeline.py +++ b/bc2/core/test_pipeline.py @@ -6,6 +6,20 @@ from .pipeline import Pipeline, PipelineConfig +def test_pipeline_usage_reporting_runtime_flag(): + ctx = Pipeline.create([]).run({"report_usage": True}) + + assert ctx.usage == { + "calls": [], + "totals": { + "calls": 0, + "estimated_cost": 0.0, + "estimated_calls": 0, + "unpriced_calls": 0, + }, + } + + def test_pipeline_simple_debug(): cfg = PipelineConfig.model_validate( { diff --git a/bc2/lib/embedding/openai.py b/bc2/lib/embedding/openai.py index 32b9d137..57d94b3a 100644 --- a/bc2/lib/embedding/openai.py +++ b/bc2/lib/embedding/openai.py @@ -6,13 +6,14 @@ from openai.types import CreateEmbeddingResponse from pydantic import BaseModel, Field, PositiveInt -from bc2.core.common.openai import OpenAIClientConfig +from bc2.core.common.openai import OpenAIClientConfig, _openai_provider from bc2.core.common.openai_metadata import ( EmbeddingModelMeta, ModelNotFound, get_embedding_model_meta, get_encoding_for_model, ) +from bc2.core.common.usage import record_usage from .base import BaseEmbeddingDriver from .embedding import Embedding @@ -171,6 +172,27 @@ def _format_result(self, result: CreateEmbeddingResponse) -> Embedding: if self.config.model_version: version += f"@{self.config.model_version}" + usage = getattr(result, "usage", None) + if usage is not None: + input_tokens = getattr(usage, "prompt_tokens", None) + total_tokens = getattr(usage, "total_tokens", None) + record_usage( + { + "provider": _openai_provider(self.client), + "service": "embeddings", + "model": self.config.openai_model or result_model, + "deployment": self.config.model, + "usage": { + k: v + for k, v in { + "input_tokens": input_tokens, + "total_tokens": total_tokens, + }.items() + if isinstance(v, int) + }, + } + ) + return Embedding( result.data[0].embedding, vendor=vendor_name, diff --git a/bc2/lib/embedding/test_openai.py b/bc2/lib/embedding/test_openai.py index 945cc3b1..cf2b7ea3 100644 --- a/bc2/lib/embedding/test_openai.py +++ b/bc2/lib/embedding/test_openai.py @@ -5,6 +5,11 @@ from bc2.core.common.openai import OpenAIClientConfig from bc2.core.common.openai_metadata import get_encoding_for_model +from bc2.core.common.usage import ( + create_usage_tracker, + usage_operation, + usage_tracking, +) from .openai import ( OpenAIEmbeddingConfig, @@ -46,6 +51,7 @@ def _mock_embedding_response(model: str = "text-embedding-3-large") -> MagicMock response = MagicMock() response.data = [MagicMock(embedding=[0.0, 0.1, 0.2])] response.model = model + response.usage = MagicMock(prompt_tokens=12, total_tokens=12) return response @@ -96,6 +102,32 @@ def test_embed_config_dimensions_override(): assert config.generator.model_dimensions == 512 +def test_embed_records_usage(): + client = MagicMock(spec=OpenAI) + client.base_url = "https://example.openai.azure.com/openai/v1/" + client.embeddings.create.return_value = _mock_embedding_response() + aclient = MagicMock(spec=AsyncOpenAI) + driver = OpenAIEmbeddingDriver( + client, + aclient, + OpenAIEmbeddingGeneratorConfig( + model="embedding-deployment", + openai_model="text-embedding-3-large", + ), + ) + created = create_usage_tracker({"report_usage": True}) + assert created is not None + report, tracker = created + + with usage_tracking(tracker), usage_operation("inspect:embed"): + driver.embed("hello") + + call = report["calls"][0] + assert call["provider"] == "azure" + assert call["operation"] == "inspect:embed" + assert call["usage"] == {"input_tokens": 12, "total_tokens": 12} + + @pytest.mark.asyncio async def test_embed_async_passes_dimensions(): client = MagicMock(spec=OpenAI) From edd0d975b7facf909f22341c017a5d79c2045896 Mon Sep 17 00:00:00 2001 From: Joe Nudell Date: Fri, 10 Jul 2026 13:05:43 -0400 Subject: [PATCH 2/3] Potential fix for pull request finding 'CodeQL / Incomplete URL substring sanitization' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- bc2/core/common/openai.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/bc2/core/common/openai.py b/bc2/core/common/openai.py index 679fbf29..6e955a4d 100644 --- a/bc2/core/common/openai.py +++ b/bc2/core/common/openai.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from functools import cached_property from typing import Any, Generic, Literal, Sequence, Type, TypeVar, cast +from urllib.parse import urlparse from openai import AsyncOpenAI, OpenAI from openai.types.responses import ( @@ -606,7 +607,9 @@ def _record_response_usage( def _openai_provider(client: OpenAI | AsyncOpenAI) -> str: base_url = str(getattr(client, "base_url", "")) - if "/openai/" in base_url or ".openai.azure.com" in base_url: + parsed = urlparse(base_url) + host = (parsed.hostname or "").lower() + if "/openai/" in parsed.path or host == "openai.azure.com" or host.endswith(".openai.azure.com"): return "azure" return "openai" From 4c8cb55c7b099635ad4c3520b1283087071e69ab Mon Sep 17 00:00:00 2001 From: Joe Nudell Date: Fri, 10 Jul 2026 13:08:00 -0400 Subject: [PATCH 3/3] fix formatting --- bc2/core/common/openai.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/bc2/core/common/openai.py b/bc2/core/common/openai.py index 6e955a4d..ee9d2806 100644 --- a/bc2/core/common/openai.py +++ b/bc2/core/common/openai.py @@ -609,7 +609,11 @@ def _openai_provider(client: OpenAI | AsyncOpenAI) -> str: base_url = str(getattr(client, "base_url", "")) parsed = urlparse(base_url) host = (parsed.hostname or "").lower() - if "/openai/" in parsed.path or host == "openai.azure.com" or host.endswith(".openai.azure.com"): + if ( + "/openai/" in parsed.path + or host == "openai.azure.com" + or host.endswith(".openai.azure.com") + ): return "azure" return "openai"