Run async batching in Outlines engine when native batching not available (#170)

* feat: Use async batching for models without batch support in outlines.

* feat: Use async batching for models without batch support in outlines.

* fix: Expose OPENROUTER_API_KEY.

* test: Remove redundant test.

---------

Co-authored-by: Raphael Mitsch <raphael@climatiq.com>

Author: rmitsch

.github/workflows/test.yml

@@ -97,6 +97,7 @@ jobs:
     - name: Run tests
       env:
         ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+        OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }}
       run: |
         source .venv/bin/activate
         pytest -x --cov --cov-report=xml -m "not slow"

sieves/engines/outlines_.py

@@ -73,7 +73,17 @@ def generate(prompts: list[str]) -> Iterable[Result]:
                     results = generator.batch(prompts, **self._inference_kwargs)
                 # Batch mode is not implemented for all Outlines wrappers. Fall back to single-prompt mode in that case.
                 except NotImplementedError:
-                    results = [generator(prompt, **self._inference_kwargs) for prompt in prompts]
+
+                    async def generate_async(prompt: str) -> Result | None:
+                        """Generate result async.
+
+                        :param prompt: Prompt to generate result for.
+                        :return: Result for prompt. Results are None if corresponding prompt failed.
+                        """
+                        return generator(prompt, **self._inference_kwargs)
+
+                    calls = [generate_async(prompt) for prompt in prompts]
+                    results = self._execute_async_calls(calls)
 
                 if inference_mode == InferenceMode.json:
                     assert len(results) == len(prompts)

sieves/tests/tasks/engines/test_edge_cases.py

@@ -0,0 +1,32 @@
+"""Test engine-related edge cases, e.g. specific model-task combinations or engine behavior under certain conditions."""
+import os
+
+import openai
+import outlines
+
+from sieves import tasks, GenerationSettings, Pipeline, Doc
+
+
+def test_openai_outlines_batching() -> None:
+    """Test Outlines batching fallback uses async batching."""
+    model = "gpt-5-nano"
+    client = openai.OpenAI(
+        base_url="https://openrouter.ai/api/v1",
+        api_key=os.environ["OPENROUTER_API_KEY"],
+    )
+    model = outlines.from_openai(client, model_name=model)
+
+    classifier = tasks.Classification(
+        task_id="procedures_classifier",
+        labels=["Fruit", "Vegetable"],
+        model=model,
+        generation_settings=GenerationSettings(
+            batch_size=10,
+            strict_mode=False,
+            inference_kwargs={"max_tokens": 200},
+        ),
+        multi_label=True,
+    )
+
+    docs = [Doc(text="Apple"), Doc(text="Carrot"), Doc(text="Watermelon")]
+    list(Pipeline(classifier)(docs))

sieves/tests/tasks/postprocessing/test_distillation.py

@@ -199,32 +199,3 @@ def test_serialization(classification_docs, batch_runtime) -> None:
 
 
     Pipeline.deserialize(config=config, tasks_kwargs=[{"model": batch_runtime.model}, {}])
-
-
-@pytest.mark.skip(reason="No OpenAI API key available in GitHub CI yet")
-def test_distillation_with_openai_model() -> None:
-    """Test distillation with an OpenAI model.
-
-    See https://github.com/MantisAI/sieves/issues/162.
-    """
-    client = openai.OpenAI(api_key=os.environ['OPENAI_API_KEY'])
-    model = outlines.from_openai(client=client, model_name='gpt-5-nano')
-
-    classifier = Classification(
-        task_id='classifier',
-        labels=['Fruit', 'Vegetable'],
-        model=model,
-        generation_settings=GenerationSettings(batch_size=32, strict_mode=False),
-    )
-
-    distiller = Distillation(
-        target_task_id='classifier',
-        base_model_id='sentence-transformers/paraphrase-mpnet-base-v2',
-        framework=DistillationFramework.setfit,
-        output_path='./model',
-        train_frac=.5,
-        val_frac=.5,
-    )
-
-    pipe = classifier + distiller
-    list(pipe([Doc(text='Apple'), Doc(text='Cucumber'), Doc(text="Broccoli"), Doc(text='Pomegranate')]))