97e91a83f3
Ruff / Ruff (push) Waiting to run
Test / Core Tests (push) Waiting to run
Test / Offline Coverage Tests (Python 3.10) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.11) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.12) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.13) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.9) (push) Waiting to run
Test / Full Coverage (Python 3.11) (push) Waiting to run
Test / Core Provider Tests (OpenAI) (push) Blocked by required conditions
Test / Core Provider Tests (Anthropic) (push) Blocked by required conditions
Test / Core Provider Tests (Google) (push) Blocked by required conditions
Test / Core Provider Tests (Other) (push) Blocked by required conditions
Test / Anthropic Tests (push) Blocked by required conditions
Test / Gemini Tests (push) Blocked by required conditions
Test / Google GenAI Tests (push) Blocked by required conditions
Test / Vertex AI Tests (push) Blocked by required conditions
Test / OpenAI Tests (push) Blocked by required conditions
Test / Writer Tests (push) Blocked by required conditions
Test / Auto Client Tests (push) Blocked by required conditions
ty / type-check (push) Waiting to run
530 lines
18 KiB
Python
530 lines
18 KiB
Python
"""Behavior tests for the OpenAI batch provider and the provider contract."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import importlib.util
|
|
import runpy
|
|
from collections.abc import Sequence
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import openai
|
|
import pytest
|
|
|
|
import instructor.batch.providers as providers
|
|
from instructor.batch.models import BatchJobInfo, BatchStatus
|
|
from instructor.batch.providers.base import BatchProvider
|
|
from instructor.batch.providers.openai import OpenAIProvider
|
|
|
|
pytestmark = pytest.mark.unit
|
|
|
|
|
|
class BatchResponse:
|
|
"""Small SDK-shaped batch response used by the provider tests."""
|
|
|
|
def __init__(
|
|
self,
|
|
batch_id: str = "batch_123",
|
|
status: str = "completed",
|
|
output_file_id: str | None = "file_output",
|
|
request_counts: Any = None,
|
|
) -> None:
|
|
self.id = batch_id
|
|
self.status = status
|
|
self.output_file_id = output_file_id
|
|
self.request_counts = request_counts
|
|
self.created_at = 1_700_000_000
|
|
|
|
def model_dump(self) -> dict[str, Any]:
|
|
counts = self.request_counts
|
|
return {
|
|
"id": self.id,
|
|
"status": self.status,
|
|
"created_at": self.created_at,
|
|
"request_counts": {
|
|
"total": getattr(counts, "total", None),
|
|
"completed": getattr(counts, "completed", None),
|
|
"failed": getattr(counts, "failed", None),
|
|
},
|
|
"input_file_id": "file_input",
|
|
"output_file_id": self.output_file_id,
|
|
"metadata": {"source": "unit-test"},
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
}
|
|
|
|
|
|
class FilesAPI:
|
|
def __init__(self, text: str = '{"ok": true}\n') -> None:
|
|
self.text = text
|
|
self.uploads: list[tuple[bytes, str]] = []
|
|
self.content_ids: list[str] = []
|
|
self.create_error: Exception | None = None
|
|
self.content_error: Exception | None = None
|
|
|
|
def create(self, *, file: Any, purpose: str) -> SimpleNamespace:
|
|
if self.create_error:
|
|
raise self.create_error
|
|
self.uploads.append((file.read(), purpose))
|
|
return SimpleNamespace(id="file_input")
|
|
|
|
def content(self, file_id: str) -> SimpleNamespace:
|
|
if self.content_error:
|
|
raise self.content_error
|
|
self.content_ids.append(file_id)
|
|
return SimpleNamespace(text=self.text)
|
|
|
|
|
|
class BatchesAPI:
|
|
def __init__(self, responses: Sequence[BatchResponse] | None = None) -> None:
|
|
self.responses = responses or [BatchResponse()]
|
|
self.retrieve_ids: list[str] = []
|
|
self.created: list[dict[str, Any]] = []
|
|
self.cancelled: list[str] = []
|
|
self.limits: list[int] = []
|
|
self.create_error: Exception | None = None
|
|
self.retrieve_error: Exception | None = None
|
|
self.cancel_error: Exception | None = None
|
|
self.list_error: Exception | None = None
|
|
|
|
def create(self, **kwargs: Any) -> SimpleNamespace:
|
|
if self.create_error:
|
|
raise self.create_error
|
|
self.created.append(kwargs)
|
|
return SimpleNamespace(id="batch_created")
|
|
|
|
def retrieve(self, batch_id: str) -> BatchResponse:
|
|
if self.retrieve_error:
|
|
raise self.retrieve_error
|
|
self.retrieve_ids.append(batch_id)
|
|
index = min(len(self.retrieve_ids) - 1, len(self.responses) - 1)
|
|
return self.responses[index]
|
|
|
|
def cancel(self, batch_id: str) -> BatchResponse:
|
|
if self.cancel_error:
|
|
raise self.cancel_error
|
|
self.cancelled.append(batch_id)
|
|
return BatchResponse(batch_id=batch_id, status="cancelled")
|
|
|
|
def list(self, *, limit: int) -> SimpleNamespace:
|
|
if self.list_error:
|
|
raise self.list_error
|
|
self.limits.append(limit)
|
|
return SimpleNamespace(data=self.responses)
|
|
|
|
|
|
class OpenAIClient:
|
|
def __init__(self, responses: list[BatchResponse] | None = None) -> None:
|
|
self.files = FilesAPI()
|
|
self.batches = BatchesAPI(responses)
|
|
|
|
|
|
def install_client(
|
|
monkeypatch: pytest.MonkeyPatch, client: OpenAIClient
|
|
) -> OpenAIProvider:
|
|
monkeypatch.setattr(openai, "OpenAI", lambda: client)
|
|
return OpenAIProvider()
|
|
|
|
|
|
def test_submit_batch_uploads_file_and_uses_default_metadata(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
batch_file = tmp_path / "requests.jsonl"
|
|
batch_file.write_bytes(b'{"custom_id": "one"}\n')
|
|
client = OpenAIClient()
|
|
provider = install_client(monkeypatch, client)
|
|
|
|
batch_id = provider.submit_batch(str(batch_file))
|
|
|
|
assert batch_id == "batch_created"
|
|
assert client.files.uploads == [(b'{"custom_id": "one"}\n', "batch")]
|
|
assert client.batches.created == [
|
|
{
|
|
"input_file_id": "file_input",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"metadata": {"description": "Instructor batch job"},
|
|
}
|
|
]
|
|
|
|
|
|
def test_submit_batch_rewinds_buffer_and_passes_custom_options(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
buffer = io.BytesIO(b'{"custom_id": "two"}\n')
|
|
buffer.seek(5)
|
|
client = OpenAIClient()
|
|
provider = install_client(monkeypatch, client)
|
|
|
|
batch_id = provider.submit_batch(
|
|
buffer, metadata={"source": "daily-import"}, completion_window="48h"
|
|
)
|
|
|
|
assert batch_id == "batch_created"
|
|
assert client.files.uploads == [(b'{"custom_id": "two"}\n', "batch")]
|
|
assert client.batches.created == [
|
|
{
|
|
"input_file_id": "file_input",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "48h",
|
|
"metadata": {"source": "daily-import"},
|
|
}
|
|
]
|
|
|
|
|
|
def test_submit_batch_rejects_invalid_input_without_credentials(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
|
invalid_input: Any = 123
|
|
|
|
with pytest.raises(ValueError, match="Unsupported file_path_or_buffer type"):
|
|
OpenAIProvider().submit_batch(invalid_input)
|
|
|
|
|
|
@pytest.mark.parametrize("error", [ValueError("bad input"), TypeError("bad file")])
|
|
def test_submit_batch_preserves_validation_errors(
|
|
monkeypatch: pytest.MonkeyPatch, error: Exception
|
|
) -> None:
|
|
client = OpenAIClient()
|
|
client.files.create_error = error
|
|
provider = install_client(monkeypatch, client)
|
|
|
|
with pytest.raises(type(error), match=str(error)) as caught:
|
|
provider.submit_batch(io.BytesIO(b"{}\n"))
|
|
|
|
assert caught.value is error
|
|
assert client.batches.created == []
|
|
|
|
|
|
def test_submit_batch_wraps_service_errors(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
client = OpenAIClient()
|
|
client.batches.create_error = ConnectionError("service unavailable")
|
|
provider = install_client(monkeypatch, client)
|
|
|
|
with pytest.raises(RuntimeError, match="Failed to submit OpenAI batch") as caught:
|
|
provider.submit_batch(io.BytesIO(b"{}\n"))
|
|
|
|
assert isinstance(caught.value.__cause__, ConnectionError)
|
|
assert "service unavailable" in str(caught.value)
|
|
|
|
|
|
def test_get_status_maps_sdk_response_and_missing_count_fields(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
batch = BatchResponse(status="in_progress", request_counts=SimpleNamespace(total=4))
|
|
client = OpenAIClient([batch])
|
|
provider = install_client(monkeypatch, client)
|
|
|
|
status = provider.get_status("batch_123")
|
|
|
|
assert status == {
|
|
"id": "batch_123",
|
|
"status": "in_progress",
|
|
"created_at": 1_700_000_000,
|
|
"request_counts": {"total": 4, "completed": 0, "failed": 0},
|
|
}
|
|
assert client.batches.retrieve_ids == ["batch_123"]
|
|
|
|
|
|
def test_get_status_wraps_sdk_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
client = OpenAIClient()
|
|
client.batches.retrieve_error = ConnectionError("status service unavailable")
|
|
provider = install_client(monkeypatch, client)
|
|
|
|
with pytest.raises(Exception, match="Failed to get OpenAI batch status") as caught:
|
|
provider.get_status("batch_missing")
|
|
|
|
assert isinstance(caught.value.__cause__, ConnectionError)
|
|
|
|
|
|
@pytest.mark.parametrize("operation", ["retrieve", "download"])
|
|
def test_completed_batch_reads_or_writes_results(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, operation: str
|
|
) -> None:
|
|
counts = SimpleNamespace(total=2, completed=1, failed=1)
|
|
batch = BatchResponse(request_counts=counts)
|
|
client = OpenAIClient([batch])
|
|
client.files.text = '{"custom_id": "one", "response": {}}\n'
|
|
provider = install_client(monkeypatch, client)
|
|
destination = tmp_path / "results.jsonl"
|
|
|
|
if operation == "retrieve":
|
|
result = provider.retrieve_results("batch_123")
|
|
assert result == client.files.text
|
|
else:
|
|
result = provider.download_results("batch_123", str(destination))
|
|
assert result is None
|
|
assert destination.read_text() == client.files.text
|
|
|
|
assert client.batches.retrieve_ids == ["batch_123"]
|
|
assert client.files.content_ids == ["file_output"]
|
|
|
|
|
|
@pytest.mark.parametrize("operation", ["retrieve", "download"])
|
|
def test_results_wait_for_output_file_then_succeed(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
capsys: pytest.CaptureFixture[str],
|
|
operation: str,
|
|
) -> None:
|
|
waiting = BatchResponse(output_file_id=None)
|
|
ready = BatchResponse(output_file_id="file_ready")
|
|
client = OpenAIClient([waiting, ready])
|
|
provider = install_client(monkeypatch, client)
|
|
waits: list[int] = []
|
|
monkeypatch.setattr("time.sleep", waits.append)
|
|
destination = tmp_path / "results.jsonl"
|
|
|
|
if operation == "retrieve":
|
|
assert provider.retrieve_results("batch_123") == client.files.text
|
|
else:
|
|
provider.download_results("batch_123", str(destination))
|
|
assert destination.read_text() == client.files.text
|
|
|
|
output = capsys.readouterr().out
|
|
assert waits == [5]
|
|
assert "waiting 5s (attempt 1/10)" in output
|
|
assert "Output file now available: file_ready" in output
|
|
assert client.batches.retrieve_ids == ["batch_123", "batch_123"]
|
|
assert client.files.content_ids == ["file_ready"]
|
|
|
|
|
|
@pytest.mark.parametrize("operation", ["retrieve", "download"])
|
|
def test_results_reject_noncompleted_and_all_failed_batches(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, operation: str
|
|
) -> None:
|
|
pending_client = OpenAIClient([BatchResponse(status="in_progress")])
|
|
provider = install_client(monkeypatch, pending_client)
|
|
destination = tmp_path / "results.jsonl"
|
|
|
|
with pytest.raises(Exception, match="Batch not completed, status: in_progress"):
|
|
if operation == "retrieve":
|
|
provider.retrieve_results("batch_pending")
|
|
else:
|
|
provider.download_results("batch_pending", str(destination))
|
|
|
|
failed_counts = SimpleNamespace(total=3, completed=0, failed=3)
|
|
failed_client = OpenAIClient([BatchResponse(request_counts=failed_counts)])
|
|
provider = install_client(monkeypatch, failed_client)
|
|
with pytest.raises(Exception, match="All 3 batch requests failed"):
|
|
if operation == "retrieve":
|
|
provider.retrieve_results("batch_failed")
|
|
else:
|
|
provider.download_results("batch_failed", str(destination))
|
|
|
|
assert pending_client.files.content_ids == []
|
|
assert failed_client.files.content_ids == []
|
|
assert not destination.exists()
|
|
|
|
|
|
@pytest.mark.parametrize("operation", ["retrieve", "download"])
|
|
def test_results_stop_when_status_changes_while_waiting(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, operation: str
|
|
) -> None:
|
|
responses = [
|
|
BatchResponse(output_file_id=None),
|
|
BatchResponse(status="failed", output_file_id=None),
|
|
]
|
|
client = OpenAIClient(responses)
|
|
provider = install_client(monkeypatch, client)
|
|
waits: list[int] = []
|
|
monkeypatch.setattr("time.sleep", waits.append)
|
|
destination = tmp_path / "results.jsonl"
|
|
|
|
with pytest.raises(Exception, match="Batch status changed to failed"):
|
|
if operation == "retrieve":
|
|
provider.retrieve_results("batch_123")
|
|
else:
|
|
provider.download_results("batch_123", str(destination))
|
|
|
|
assert waits == [5]
|
|
assert len(client.batches.retrieve_ids) == 2
|
|
assert client.files.content_ids == []
|
|
assert not destination.exists()
|
|
|
|
|
|
@pytest.mark.parametrize("operation", ["retrieve", "download"])
|
|
def test_results_report_exhausted_output_retries(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, operation: str
|
|
) -> None:
|
|
client = OpenAIClient([BatchResponse(output_file_id=None)])
|
|
provider = install_client(monkeypatch, client)
|
|
waits: list[int] = []
|
|
monkeypatch.setattr("time.sleep", waits.append)
|
|
destination = tmp_path / "results.jsonl"
|
|
|
|
with pytest.raises(Exception, match="No output file available after 10 retries"):
|
|
if operation == "retrieve":
|
|
provider.retrieve_results("batch_123")
|
|
else:
|
|
provider.download_results("batch_123", str(destination))
|
|
|
|
assert waits == list(range(5, 15))
|
|
assert len(client.batches.retrieve_ids) == 11
|
|
assert client.files.content_ids == []
|
|
assert not destination.exists()
|
|
|
|
|
|
@pytest.mark.parametrize("operation", ["retrieve", "download"])
|
|
def test_results_wrap_file_read_errors(
|
|
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, operation: str
|
|
) -> None:
|
|
client = OpenAIClient()
|
|
client.files.content_error = ConnectionError("file service unavailable")
|
|
provider = install_client(monkeypatch, client)
|
|
destination = tmp_path / "results.jsonl"
|
|
|
|
expected = (
|
|
"Failed to retrieve OpenAI results"
|
|
if operation == "retrieve"
|
|
else "Failed to download OpenAI results"
|
|
)
|
|
with pytest.raises(Exception, match=expected) as caught:
|
|
if operation == "retrieve":
|
|
provider.retrieve_results("batch_123")
|
|
else:
|
|
provider.download_results("batch_123", str(destination))
|
|
|
|
assert isinstance(caught.value.__cause__, ConnectionError)
|
|
assert not destination.exists()
|
|
|
|
|
|
def test_cancel_delete_and_list_batches_map_sdk_results(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
counts = SimpleNamespace(total=4, completed=4, failed=0)
|
|
batch = BatchResponse(request_counts=counts)
|
|
client = OpenAIClient([batch])
|
|
provider = install_client(monkeypatch, client)
|
|
|
|
cancelled = provider.cancel_batch("batch_cancel")
|
|
deleted = provider.delete_batch("batch_123")
|
|
listed = provider.list_batches(limit=3)
|
|
|
|
assert cancelled["id"] == "batch_cancel"
|
|
assert cancelled["status"] == "cancelled"
|
|
assert client.batches.cancelled == ["batch_cancel"]
|
|
assert deleted == {
|
|
"id": "batch_123",
|
|
"status": "completed",
|
|
"message": "OpenAI does not support batch deletion",
|
|
}
|
|
assert client.batches.limits == [3]
|
|
assert len(listed) == 1
|
|
assert isinstance(listed[0], BatchJobInfo)
|
|
assert listed[0].provider == "openai"
|
|
assert listed[0].status == BatchStatus.COMPLETED
|
|
assert listed[0].request_counts.total == 4
|
|
assert listed[0].files.output_file_id == "file_output"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("operation", "error_attribute", "expected"),
|
|
[
|
|
("cancel", "cancel_error", "Failed to cancel OpenAI batch"),
|
|
("delete", "retrieve_error", "Failed to delete OpenAI batch"),
|
|
("list", "list_error", "Failed to list OpenAI batches"),
|
|
],
|
|
)
|
|
def test_batch_actions_wrap_sdk_errors(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
operation: str,
|
|
error_attribute: str,
|
|
expected: str,
|
|
) -> None:
|
|
client = OpenAIClient()
|
|
setattr(
|
|
client.batches, error_attribute, ConnectionError("batch service unavailable")
|
|
)
|
|
provider = install_client(monkeypatch, client)
|
|
|
|
with pytest.raises(Exception, match=expected) as caught:
|
|
if operation == "cancel":
|
|
provider.cancel_batch("batch_123")
|
|
elif operation == "delete":
|
|
provider.delete_batch("batch_123")
|
|
else:
|
|
provider.list_batches()
|
|
|
|
assert isinstance(caught.value.__cause__, ConnectionError)
|
|
|
|
|
|
def test_provider_factory_selects_available_providers() -> None:
|
|
openai_provider = providers.get_provider("openai")
|
|
anthropic_provider = providers.get_provider("anthropic")
|
|
|
|
openai_provider_type = providers.OpenAIProvider
|
|
anthropic_provider_type = providers.AnthropicProvider
|
|
assert openai_provider_type is not None
|
|
assert anthropic_provider_type is not None
|
|
assert isinstance(openai_provider, openai_provider_type)
|
|
assert isinstance(anthropic_provider, anthropic_provider_type)
|
|
assert isinstance(openai_provider, BatchProvider)
|
|
assert isinstance(anthropic_provider, BatchProvider)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("provider_name", "attribute", "expected"),
|
|
[
|
|
("openai", "OpenAIProvider", "OpenAI is not installed"),
|
|
("anthropic", "AnthropicProvider", "Anthropic is not installed"),
|
|
],
|
|
)
|
|
def test_provider_factory_reports_missing_optional_provider(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
provider_name: str,
|
|
attribute: str,
|
|
expected: str,
|
|
) -> None:
|
|
monkeypatch.setattr(providers, attribute, None)
|
|
|
|
with pytest.raises(ValueError, match=expected):
|
|
providers.get_provider(provider_name)
|
|
|
|
|
|
def test_provider_factory_rejects_unknown_provider() -> None:
|
|
with pytest.raises(ValueError, match="Unsupported provider: unknown"):
|
|
providers.get_provider("unknown")
|
|
|
|
|
|
def test_provider_module_handles_missing_optional_sdks(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
original_find_spec = importlib.util.find_spec
|
|
monkeypatch.setattr(
|
|
importlib.util,
|
|
"find_spec",
|
|
lambda name: None
|
|
if name in {"openai", "anthropic"}
|
|
else original_find_spec(name),
|
|
)
|
|
|
|
namespace = runpy.run_path(
|
|
providers.__file__, run_name="instructor.batch.providers.without_sdks"
|
|
)
|
|
|
|
assert namespace["OpenAIProvider"] is None
|
|
assert namespace["AnthropicProvider"] is None
|
|
with pytest.raises(ValueError, match="OpenAI is not installed"):
|
|
namespace["get_provider"]("openai")
|
|
with pytest.raises(ValueError, match="Anthropic is not installed"):
|
|
namespace["get_provider"]("anthropic")
|
|
|
|
|
|
def test_batch_provider_contract_requires_and_defines_all_operations() -> None:
|
|
expected = {
|
|
"submit_batch",
|
|
"get_status",
|
|
"retrieve_results",
|
|
"download_results",
|
|
"cancel_batch",
|
|
"delete_batch",
|
|
"list_batches",
|
|
}
|
|
assert BatchProvider.__abstractmethods__ == expected
|
|
with pytest.raises(TypeError, match="abstract methods"):
|
|
BatchProvider()
|