feat: add native Azure OpenAI batch backend (#538) (#727)

Adds a native `backend="azure"` batch request processor, since Azure OpenAI
was previously only reachable via the litellm backend, which doesn't support
batch mode. Batch was already supported natively for OpenAI, Anthropic,
Gemini, and Mistral, but not Azure.

Azure's OpenAI-compatible "v1" API surface (https://<resource>.openai.azure.com/openai/v1/)
lets AzureBatchRequestProcessor subclass OpenAIBatchRequestProcessor directly and
reuse its file/batch lifecycle handling almost unmodified, rather than
reimplementing it. It accounts for two documented differences from OpenAI's
own API: the batches.create() `endpoint` argument omits the "/v1" prefix, and
the request body's "model" field must be the Azure deployment name rather
than curator's `model` config (which stays the underlying model name for
litellm-based cost lookups).

Also adds a `_AzureCostProcessor` (mirroring the existing KlusterAI/
InferenceNet pattern) to prefer litellm's `azure/`-prefixed pricing table,
and a small refactor extracting OpenAIBatchRequestProcessor's hardcoded batch
endpoint string into an overridable class attribute so Azure doesn't need to
duplicate create_batch().

Not verified against a live Azure resource (no Azure credentials available
in this environment) - the design follows Microsoft's current official
documentation, and the integration test is left @pytest.mark.skip()-decorated
pending a run against a real resource.

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

Author: EpicSRN021

src/bespokelabs/curator/cost.py

@@ -125,9 +125,30 @@ def cost(self, *, completion_window="*", **kwargs):
         return super().cost(completion_window=completion_window, **kwargs) * times
 
 
+class _AzureCostProcessor(_LitellmCostProcessor):
+    """Cost processor for Azure OpenAI.
+
+    `config.model` stays the underlying model name (e.g. "gpt-4o") rather than the Azure
+    deployment name, since deployment names are user-chosen and aren't in litellm's cost table.
+    This prefers litellm's `azure/`-prefixed pricing (which can differ slightly from the bare
+    OpenAI entry) when available, falling back to the default OpenAI-style lookup otherwise.
+    """
+
+    def cost(self, *, completion_window="*", **kwargs):
+        azure_model = f"azure/{self.config.model}"
+        if azure_model in litellm.model_cost:
+            kwargs["model"] = azure_model
+            cost_to_complete = litellm.completion_cost(**kwargs)
+            if self.batch:
+                cost_to_complete *= 0.5
+            return cost_to_complete
+        return super().cost(completion_window=completion_window, **kwargs)
+
+
 COST_PROCESSOR = defaultdict(lambda: _LitellmCostProcessor)
 COST_PROCESSOR["klusterai"] = _KlusterAICostProcessor
 COST_PROCESSOR["inference.net"] = _InferenceNetCostProcessor
+COST_PROCESSOR["azure"] = _AzureCostProcessor
 
 
 def cost_processor_factory(backend, config=None, batch=False):

src/bespokelabs/curator/request_processor/_factory.py

@@ -177,4 +177,14 @@ def create(
 
             return MistralBatchRequestProcessor(config)
 
+        if backend == "azure":
+            config.api_key = config.api_key or os.getenv("AZURE_OPENAI_API_KEY")
+            if not config.api_key:
+                raise ValueError("AZURE_OPENAI_API_KEY is not set.")
+            if not batch:
+                raise ValueError("Only batch mode is supported with the Azure backend.")
+            from bespokelabs.curator.request_processor.batch.azure_batch_request_processor import AzureBatchRequestProcessor
+
+            return AzureBatchRequestProcessor(config)
+
         raise ValueError(f"Unknown backend: {backend}")

src/bespokelabs/curator/request_processor/batch/azure_batch_request_processor.py

@@ -0,0 +1,53 @@
+from litellm import model_cost
+
+from bespokelabs.curator.request_processor.batch.openai_batch_request_processor import OpenAIBatchRequestProcessor
+from bespokelabs.curator.request_processor.config import BatchRequestProcessorConfig
+from bespokelabs.curator.types.generic_request import GenericRequest
+
+
+class AzureBatchRequestProcessor(OpenAIBatchRequestProcessor):
+    """Azure OpenAI-specific implementation of the BatchRequestProcessor.
+
+    Uses Azure OpenAI's OpenAI-SDK-compatible `v1` API surface
+    (`https://<resource>.openai.azure.com/openai/v1/`), so file/batch lifecycle handling is
+    inherited from OpenAIBatchRequestProcessor unmodified. Azure diverges from OpenAI in two ways
+    this class accounts for:
+
+    - `client.batches.create()`'s `endpoint` argument omits the "/v1" prefix ("/chat/completions",
+      not "/v1/chat/completions"), even though the batch file's per-line "url" keeps it.
+    - the request body's "model" field must be set to the Azure deployment name, not curator's
+      `model` config (which stays the underlying model name, e.g. "gpt-4o", used for cost lookups).
+    """
+
+    _BATCH_ENDPOINT = "/chat/completions"
+
+    def __init__(self, config: BatchRequestProcessorConfig) -> None:
+        """Initialize the AzureBatchRequestProcessor."""
+        if not config.base_url:
+            raise ValueError("Azure OpenAI backend requires `base_url` set to your resource's v1 endpoint, e.g. https://<resource>.openai.azure.com/openai/v1/")
+        if not config.azure_deployment:
+            raise ValueError("Azure OpenAI backend requires `azure_deployment` set to your batch deployment name.")
+
+        super().__init__(config, compatible_provider="azure")
+        self.web_dashboard = "https://oai.azure.com/resource/batch"
+
+    @property
+    def backend(self):
+        """Backend property."""
+        return "azure"
+
+    def create_api_specific_request_batch(self, generic_request: GenericRequest) -> dict:
+        """Creates an Azure-specific request body, targeting the configured deployment name."""
+        request = super().create_api_specific_request_batch(generic_request)
+        request["body"]["model"] = self.config.azure_deployment
+        return request
+
+    def set_model_cost(self):
+        """Set tracker cost information, preferring litellm's `azure/`-prefixed pricing table."""
+        azure_model = f"azure/{self.prompt_formatter.model_name}"
+        if azure_model in model_cost:
+            self.tracker.input_cost_per_million = (model_cost[azure_model]["input_cost_per_token"] * 1_000_000) * 0.5
+            self.tracker.output_cost_per_million = (model_cost[azure_model]["output_cost_per_token"] * 1_000_000) * 0.5
+            self.tracker.start_tracker(self._tracker_console)
+        else:
+            super().set_model_cost()

src/bespokelabs/curator/request_processor/batch/openai_batch_request_processor.py

@@ -31,6 +31,10 @@ class OpenAIBatchRequestProcessor(BaseBatchRequestProcessor, OpenAIRequestMixin)
     file uploads, batch submissions, and result retrieval.
     """
 
+    # Overridable so OpenAI-compatible providers with a slightly different batch endpoint
+    # convention (e.g. Azure's, which omits the "/v1" prefix) don't need to duplicate create_batch.
+    _BATCH_ENDPOINT = "/v1/chat/completions"
+
     def __init__(self, config: BatchRequestProcessorConfig, compatible_provider=None) -> None:
         """Initialize the OpenAIBatchRequestProcessor."""
         super().__init__(config)
@@ -295,7 +299,7 @@ async def create_batch(self, batch_file_id: str, metadata: dict) -> Batch:
         try:
             batch = await self.client.batches.create(
                 input_file_id=batch_file_id,
-                endpoint="/v1/chat/completions",
+                endpoint=self._BATCH_ENDPOINT,
                 completion_window=self.config.completion_window,
                 metadata=metadata,
             )

src/bespokelabs/curator/request_processor/config.py

@@ -70,13 +70,18 @@ class BatchRequestProcessorConfig(RequestProcessorConfig):
         delete_successful_batch_files: Whether to delete batch files after successful processing
         delete_failed_batch_files: Whether to delete batch files after failed processing
         completion_window: Time window to wait for batch completion
+        azure_deployment: Azure OpenAI deployment name to target (required for the Azure backend).
+            This is distinct from `model`, which stays the underlying model name (e.g. "gpt-4o")
+            used for litellm-based cost lookups, since Azure deployment names are user-chosen and
+            don't map 1:1 to a model name.
     """
 
     batch_size: t.Union[int, str] = Field(default=10_000)
     batch_check_interval: int = Field(default=60, gt=0)
     delete_successful_batch_files: bool = False
     delete_failed_batch_files: bool = False
     completion_window: str = "24h"
+    azure_deployment: str | None = None
 
     def __post_init__(self):
         """Post-initialization hook to validate batch size."""
@@ -165,6 +170,7 @@ class BatchBackendParams(BaseBackendParams, total=False):
     batch_check_interval: t.Optional[int]
     delete_successful_batch_files: t.Optional[bool]
     delete_failed_batch_files: t.Optional[bool]
+    azure_deployment: t.Optional[str]
 
 
 class OfflineBackendParams(BaseBackendParams, total=False):

tests/integrations/test_azure_batch.py

@@ -0,0 +1,53 @@
+import os
+
+import pytest
+from dotenv import load_dotenv
+
+from bespokelabs import curator
+
+load_dotenv(dotenv_path=".env", verbose=True)
+
+
+def generate_test_prompts():
+    """Generate a batch of 4 test prompts to test batch processing."""
+    return [
+        "Describe the impact of quantum computing on AI in 2 sentences.",
+        "Give a creative business idea for a startup in 2025 in 1 sentence.",
+        "What is 2*10?",
+        "Who is the current prime-minister of India?",
+    ]
+
+
+@pytest.mark.skip()
+def test_azure_batch_response():
+    """Test sending a batch of requests to an Azure OpenAI deployment via Curator.
+
+    Requires AZURE_OPENAI_API_KEY, AZURE_OPENAI_BASE_URL (your resource's v1 endpoint, e.g.
+    https://<resource>.openai.azure.com/openai/v1/), and AZURE_OPENAI_DEPLOYMENT to be set.
+    """
+    llm = curator.LLM(
+        model_name="gpt-4o",
+        backend="azure",
+        batch=True,
+        backend_params={
+            "base_url": os.environ["AZURE_OPENAI_BASE_URL"],
+            "azure_deployment": os.environ["AZURE_OPENAI_DEPLOYMENT"],
+        },
+    )
+
+    prompts = generate_test_prompts()
+    responses = llm(prompts).dataset
+
+    df_responses = responses.to_pandas()
+
+    for i, prompt in enumerate(prompts):
+        response_text = df_responses.iloc[i]["response"]
+        assert isinstance(response_text, str), f"Response for prompt {i} should be a string"
+        assert len(response_text.strip()) > 0, f"Response for prompt {i} should not be empty"
+        print(f"\nPrompt: {prompt}\nResponse: {response_text}\n{'-' * 50}")
+
+    print("\nBatch job completed successfully.")
+
+
+if __name__ == "__main__":
+    test_azure_batch_response()

tests/unittests/test_batch.py

@@ -6,6 +6,7 @@
 from pydantic import BaseModel, Field
 
 from bespokelabs import curator
+from bespokelabs.curator.request_processor.batch.azure_batch_request_processor import AzureBatchRequestProcessor
 from bespokelabs.curator.request_processor.batch.mistral_batch_request_processor import MistralBatchRequestProcessor
 from bespokelabs.curator.request_processor.config import BatchRequestProcessorConfig
 from bespokelabs.curator.types.generic_batch import GenericBatch, GenericBatchRequestCounts, GenericBatchStatus
@@ -84,6 +85,83 @@ def test_mistral_batch_response_includes_normalized_finish_reason() -> None:
     assert response.finish_reason == "length"
 
 
+def test_azure_batch_requires_base_url() -> None:
+    """Azure batch backend needs a resource-specific base_url; there's no sensible default."""
+    config = BatchRequestProcessorConfig(model="gpt-4o", azure_deployment="my-gpt4o-deployment")
+    with pytest.raises(ValueError, match="base_url"):
+        AzureBatchRequestProcessor(config)
+
+
+def test_azure_batch_requires_deployment_name() -> None:
+    """Azure batch backend needs an explicit deployment name distinct from the model name."""
+    config = BatchRequestProcessorConfig(model="gpt-4o", base_url="https://my-resource.openai.azure.com/openai/v1/")
+    with pytest.raises(ValueError, match="azure_deployment"):
+        AzureBatchRequestProcessor(config)
+
+
+def test_azure_batch_endpoint_omits_v1_prefix() -> None:
+    """Azure's batches.create() `endpoint` argument omits the "/v1" prefix OpenAI's API uses."""
+    assert AzureBatchRequestProcessor._BATCH_ENDPOINT == "/chat/completions"
+
+
+def test_azure_batch_request_uses_deployment_name_not_model() -> None:
+    """The request body's `model` field must be the Azure deployment name, not curator's `model` config.
+
+    `model` stays the underlying model name (e.g. "gpt-4o") since that's what's used for litellm-based
+    cost lookups, while Azure requires the deployment name to route the actual API request.
+    """
+    processor = AzureBatchRequestProcessor.__new__(AzureBatchRequestProcessor)
+    processor.config = BatchRequestProcessorConfig(
+        model="gpt-4o",
+        base_url="https://my-resource.openai.azure.com/openai/v1/",
+        azure_deployment="my-gpt4o-deployment",
+    )
+
+    api_request = processor.create_api_specific_request_batch(_generic_request())
+
+    assert api_request["body"]["model"] == "my-gpt4o-deployment"
+    assert api_request["url"] == "/v1/chat/completions"
+
+
+def test_azure_cost_processor_prefixes_model_for_litellm_lookup(monkeypatch) -> None:
+    """_AzureCostProcessor should query litellm's `azure/`-prefixed pricing table, not the bare model name."""
+    import litellm
+
+    from bespokelabs.curator.cost import _AzureCostProcessor
+
+    config = BatchRequestProcessorConfig(model="gpt-4o", base_url="https://x.openai.azure.com/openai/v1/", azure_deployment="d")
+    processor = _AzureCostProcessor(config=config, batch=False)
+
+    monkeypatch.setitem(litellm.model_cost, "azure/gpt-4o", {"input_cost_per_token": 1.0, "output_cost_per_token": 1.0})
+
+    calls = {}
+
+    def fake_completion_cost(**kwargs):
+        calls.update(kwargs)
+        return 1.23
+
+    monkeypatch.setattr(litellm, "completion_cost", fake_completion_cost)
+
+    cost = processor.cost(model="gpt-4o", prompt="hi", completion="there")
+
+    assert calls["model"] == "azure/gpt-4o"
+    assert cost == 1.23
+
+
+def test_azure_cost_processor_falls_back_without_azure_specific_pricing() -> None:
+    """When litellm has no `azure/`-prefixed entry, fall back to the default bare-model lookup."""
+    from bespokelabs.curator.cost import _AzureCostProcessor
+
+    config = BatchRequestProcessorConfig(model="some-unlisted-model", base_url="https://x.openai.azure.com/openai/v1/", azure_deployment="d")
+    processor = _AzureCostProcessor(config=config, batch=False)
+
+    # Neither "azure/some-unlisted-model" nor "some-unlisted-model" are real litellm models,
+    # so the fallback lookup should just return 0.0 rather than raising.
+    cost = processor.cost(model="some-unlisted-model", prompt="hi", completion="there")
+
+    assert cost == 0.0
+
+
 def batch_call(model_name, prompts):
     dataset: Dataset = Dataset.from_dict({"prompt": prompts})
     llm = curator.LLM(