From 0f90004b140f0616173e6a3759fcb285dd8d8129 Mon Sep 17 00:00:00 2001
From: Joe Nudell <jnudell@hks.harvard.edu>
Date: Fri, 10 Jul 2026 16:41:14 -0400
Subject: [PATCH 1/2] support detection of govcloud endpoints

---
 bc2/core/common/openai.py      |  4 ++--
 bc2/core/common/test_openai.py | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/bc2/core/common/openai.py b/bc2/core/common/openai.py
index ee9d280..f0fee76 100644
--- a/bc2/core/common/openai.py
+++ b/bc2/core/common/openai.py
@@ -611,8 +611,8 @@ def _openai_provider(client: OpenAI | AsyncOpenAI) -> str:
     host = (parsed.hostname or "").lower()
     if (
         "/openai/" in parsed.path
-        or host == "openai.azure.com"
-        or host.endswith(".openai.azure.com")
+        or host in ("openai.azure.com", "openai.azure.us")
+        or host.endswith((".openai.azure.com", ".openai.azure.us"))
     ):
         return "azure"
     return "openai"
diff --git a/bc2/core/common/test_openai.py b/bc2/core/common/test_openai.py
index ffb71b3..90d8512 100644
--- a/bc2/core/common/test_openai.py
+++ b/bc2/core/common/test_openai.py
@@ -14,10 +14,25 @@
     OpenAIChatPromptInline,
     OpenAIChatTurn,
     OpenAIClientConfig,
+    _openai_provider,
 )
 from .usage import create_usage_tracker, usage_operation, usage_tracking
 
 
+@pytest.mark.parametrize(
+    "base_url",
+    [
+        "https://example.openai.azure.com/",
+        "https://example.openai.azure.us/",
+    ],
+)
+def test_openai_provider_recognizes_azure_clouds(base_url):
+    client = MagicMock()
+    client.base_url = base_url
+
+    assert _openai_provider(client) == "azure"
+
+
 def test_fix_azure_endpoint():
     cfg = OpenAIClientConfig(
         api_key="my-api-key",

From 703c3e3c1ad79c92da571e8fd5458a66cd95059c Mon Sep 17 00:00:00 2001
From: Joe Nudell <jnudell@hks.harvard.edu>
Date: Fri, 10 Jul 2026 17:46:44 -0400
Subject: [PATCH 2/2] support cost estimates for layout model

---
 bc2/core/common/azure_pricing.py      | 11 ++++++++---
 bc2/core/common/test_azure_pricing.py | 27 +++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/bc2/core/common/azure_pricing.py b/bc2/core/common/azure_pricing.py
index b80963b..87c7ffc 100644
--- a/bc2/core/common/azure_pricing.py
+++ b/bc2/core/common/azure_pricing.py
@@ -136,7 +136,12 @@ def _estimate_document_intelligence(
     ) -> dict[str, Any]:
         model = str(call.get("model") or "")
         features = call.get("features") or []
-        if model != "prebuilt-read":
+        meter_names = {
+            "prebuilt-read": "S0 Read Pages",
+            "prebuilt-layout": "S0 Pre-built Pages",
+        }
+        meter_name = meter_names.get(model)
+        if meter_name is None:
             raise AzurePricingUnavailable(
                 f"unsupported Document Intelligence model: {model or '<unknown>'}"
             )
@@ -158,11 +163,11 @@ def _estimate_document_intelligence(
         matches = [
             item
             for item in items
-            if item.get("meterName") == "S0 Read Pages"
+            if item.get("meterName") == meter_name
             and item.get("type") == "Consumption"
             and float(item.get("tierMinimumUnits") or 0) == 0
         ]
-        meter = _select_unique_price(matches, "S0 Read Pages")
+        meter = _select_unique_price(matches, meter_name)
         component = _price_quantity(pages, meter, "pages")
         return _cost_result([component], fetched_at)
 
diff --git a/bc2/core/common/test_azure_pricing.py b/bc2/core/common/test_azure_pricing.py
index 44f073e..c971ab2 100644
--- a/bc2/core/common/test_azure_pricing.py
+++ b/bc2/core/common/test_azure_pricing.py
@@ -82,6 +82,33 @@ def test_estimate_document_intelligence_read_cost(monkeypatch):
     assert estimate["estimated_cost"] == pytest.approx(0.0375)
 
 
+def test_estimate_document_intelligence_layout_cost(monkeypatch):
+    pricing = AzureRetailPricing()
+    meter = {
+        "skuName": "S0",
+        "meterName": "S0 Pre-built Pages",
+        "retailPrice": 10.0,
+        "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-layout",
+            "features": [],
+            "usage": {"pages": 25},
+        },
+        {"azure_region": "eastus"},
+    )
+
+    assert estimate["estimated_cost"] == pytest.approx(0.25)
+
+
 def test_estimate_embedding_cost(monkeypatch):
     pricing = AzureRetailPricing()
     meter = _meter("text-embedding-3-large-glbl", 0.00013, "1K")