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
410 lines
13 KiB
Python
410 lines
13 KiB
Python
"""Focused tests for the public batch processor API and result parsing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from pydantic import BaseModel
|
|
|
|
import instructor.batch.processor as processor_module
|
|
from instructor.batch.models import BatchError, BatchJobInfo, BatchSuccess
|
|
from instructor.batch.processor import BatchProcessor
|
|
|
|
|
|
pytestmark = pytest.mark.unit
|
|
|
|
|
|
class Person(BaseModel):
|
|
name: str
|
|
age: int
|
|
|
|
|
|
class RecordingProvider:
|
|
def __init__(self, results: str = "") -> None:
|
|
self.results = results
|
|
self.calls: list[tuple[str, Any]] = []
|
|
|
|
def submit_batch(
|
|
self,
|
|
file_path_or_buffer: str | io.BytesIO,
|
|
metadata: dict[str, Any] | None = None,
|
|
**kwargs: Any,
|
|
) -> str:
|
|
self.calls.append(("submit", (file_path_or_buffer, metadata, kwargs)))
|
|
return "batch-123"
|
|
|
|
def get_status(self, batch_id: str) -> dict[str, Any]:
|
|
self.calls.append(("status", batch_id))
|
|
return {"id": batch_id, "status": "completed"}
|
|
|
|
def retrieve_results(self, batch_id: str) -> str:
|
|
self.calls.append(("retrieve", batch_id))
|
|
return self.results
|
|
|
|
def download_results(self, batch_id: str, file_path: str) -> None:
|
|
self.calls.append(("download", (batch_id, file_path)))
|
|
Path(file_path).write_text(self.results)
|
|
|
|
def cancel_batch(self, batch_id: str) -> dict[str, Any]:
|
|
self.calls.append(("cancel", batch_id))
|
|
return {"id": batch_id, "status": "cancelled"}
|
|
|
|
def delete_batch(self, batch_id: str) -> dict[str, Any]:
|
|
self.calls.append(("delete", batch_id))
|
|
return {"id": batch_id, "deleted": True}
|
|
|
|
def list_batches(self, limit: int = 10) -> list[BatchJobInfo]:
|
|
self.calls.append(("list", limit))
|
|
return [
|
|
BatchJobInfo.from_openai(
|
|
{"id": "batch-123", "status": "completed", "metadata": {}}
|
|
)
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def provider(monkeypatch: pytest.MonkeyPatch) -> RecordingProvider:
|
|
provider = RecordingProvider()
|
|
monkeypatch.setattr(processor_module, "get_provider", lambda _: provider)
|
|
return provider
|
|
|
|
|
|
def openai_result(custom_id: str, content: str) -> str:
|
|
return json.dumps(
|
|
{
|
|
"custom_id": custom_id,
|
|
"response": {"body": {"choices": [{"message": {"content": content}}]}},
|
|
}
|
|
)
|
|
|
|
|
|
def test_init_splits_provider_and_model_and_rejects_invalid_model(
|
|
provider: RecordingProvider,
|
|
) -> None:
|
|
processor = BatchProcessor("openai/gpt-4.1-mini", Person)
|
|
|
|
assert processor.provider is provider
|
|
assert processor.provider_name == "openai"
|
|
assert processor.model_name == "gpt-4.1-mini"
|
|
assert processor.response_model is Person
|
|
|
|
with pytest.raises(
|
|
ValueError, match='Model string must be in format "provider/model-name"'
|
|
):
|
|
BatchProcessor("gpt-4.1-mini", Person)
|
|
|
|
|
|
def test_create_batch_file_replaces_existing_contents_and_serializes_requests(
|
|
provider: RecordingProvider, tmp_path: Path, capsys: pytest.CaptureFixture[str]
|
|
) -> None:
|
|
del provider
|
|
batch_file = tmp_path / "requests.jsonl"
|
|
batch_file.write_text("stale request\n")
|
|
processor = BatchProcessor("openai/gpt-4.1-mini", Person)
|
|
|
|
returned = processor.create_batch_from_messages(
|
|
[
|
|
[{"role": "user", "content": "Ada is 36"}],
|
|
[{"role": "user", "content": "Lin is 28"}],
|
|
],
|
|
str(batch_file),
|
|
max_tokens=321,
|
|
temperature=0.3,
|
|
)
|
|
|
|
lines = [json.loads(line) for line in batch_file.read_text().splitlines()]
|
|
assert returned == str(batch_file)
|
|
assert [line["custom_id"] for line in lines] == ["request-0", "request-1"]
|
|
assert [line["body"]["model"] for line in lines] == [
|
|
"gpt-4.1-mini",
|
|
"gpt-4.1-mini",
|
|
]
|
|
assert lines[0]["body"]["max_tokens"] == 321
|
|
assert lines[0]["body"]["temperature"] == 0.3
|
|
assert "Created batch file" in capsys.readouterr().out
|
|
|
|
|
|
def test_create_batch_file_creates_a_new_output_file(
|
|
provider: RecordingProvider, tmp_path: Path
|
|
) -> None:
|
|
del provider
|
|
batch_file = tmp_path / "new-requests.jsonl"
|
|
processor = BatchProcessor("openai/gpt-4.1-mini", Person)
|
|
|
|
returned = processor.create_batch_from_messages(
|
|
[[{"role": "user", "content": "Ada is 36"}]], str(batch_file)
|
|
)
|
|
|
|
assert returned == str(batch_file)
|
|
request = json.loads(batch_file.read_text())
|
|
assert request["custom_id"] == "request-0"
|
|
assert request["body"]["messages"] == [{"role": "user", "content": "Ada is 36"}]
|
|
|
|
|
|
def test_create_batch_buffer_is_readable_from_start(
|
|
provider: RecordingProvider, capsys: pytest.CaptureFixture[str]
|
|
) -> None:
|
|
del provider
|
|
processor = BatchProcessor("anthropic/claude-sonnet", Person)
|
|
|
|
buffer = processor.create_batch_from_messages(
|
|
[[{"role": "user", "content": "Ada is 36"}]], max_tokens=90, temperature=0.2
|
|
)
|
|
|
|
assert isinstance(buffer, io.BytesIO)
|
|
assert buffer.tell() == 0
|
|
request = json.loads(buffer.read().decode())
|
|
assert request["custom_id"] == "request-0"
|
|
assert request["params"]["model"] == "claude-sonnet"
|
|
assert request["params"]["max_tokens"] == 90
|
|
assert request["params"]["temperature"] == 0.2
|
|
assert "Created batch buffer with 1 requests" in capsys.readouterr().out
|
|
|
|
|
|
def test_provider_operations_forward_arguments_and_parse_downloaded_results(
|
|
provider: RecordingProvider, tmp_path: Path
|
|
) -> None:
|
|
provider.results = openai_result("request-0", '{"name": "Ada", "age": 36}')
|
|
processor = BatchProcessor("openai/gpt-4.1-mini", Person)
|
|
request_buffer = io.BytesIO(b"request")
|
|
|
|
assert (
|
|
processor.submit_batch(request_buffer, completion_window="24h") == "batch-123"
|
|
)
|
|
assert provider.calls[-1] == (
|
|
"submit",
|
|
(
|
|
request_buffer,
|
|
{"description": "Instructor batch job"},
|
|
{"completion_window": "24h"},
|
|
),
|
|
)
|
|
assert processor.submit_batch("requests.jsonl", metadata={"team": "search"}) == (
|
|
"batch-123"
|
|
)
|
|
assert provider.calls[-1] == (
|
|
"submit",
|
|
("requests.jsonl", {"team": "search"}, {}),
|
|
)
|
|
assert processor.get_batch_status("batch-123") == {
|
|
"id": "batch-123",
|
|
"status": "completed",
|
|
}
|
|
jobs = processor.list_batches(limit=2)
|
|
assert len(jobs) == 1
|
|
assert jobs[0].id == "batch-123"
|
|
assert provider.calls[-1] == ("list", 2)
|
|
|
|
results_file = tmp_path / "results.jsonl"
|
|
results = processor.get_results("batch-123", str(results_file))
|
|
assert isinstance(results[0], BatchSuccess)
|
|
assert results[0].result == Person(name="Ada", age=36)
|
|
assert results_file.read_text() == provider.results
|
|
assert ("retrieve", "batch-123") in provider.calls
|
|
assert ("download", ("batch-123", str(results_file))) in provider.calls
|
|
|
|
calls_before = len(provider.calls)
|
|
in_memory_results = processor.get_results("batch-123")
|
|
assert isinstance(in_memory_results[0], BatchSuccess)
|
|
assert provider.calls[calls_before:] == [("retrieve", "batch-123")]
|
|
assert processor.cancel_batch("batch-123") == {
|
|
"id": "batch-123",
|
|
"status": "cancelled",
|
|
}
|
|
assert processor.delete_batch("batch-123") == {
|
|
"id": "batch-123",
|
|
"deleted": True,
|
|
}
|
|
assert provider.calls == [
|
|
(
|
|
"submit",
|
|
(
|
|
request_buffer,
|
|
{"description": "Instructor batch job"},
|
|
{"completion_window": "24h"},
|
|
),
|
|
),
|
|
("submit", ("requests.jsonl", {"team": "search"}, {})),
|
|
("status", "batch-123"),
|
|
("list", 2),
|
|
("retrieve", "batch-123"),
|
|
("download", ("batch-123", str(results_file))),
|
|
("retrieve", "batch-123"),
|
|
("cancel", "batch-123"),
|
|
("delete", "batch-123"),
|
|
]
|
|
|
|
|
|
def test_openai_results_distinguish_success_validation_extraction_and_json_errors(
|
|
provider: RecordingProvider,
|
|
) -> None:
|
|
del provider
|
|
processor = BatchProcessor("openai/gpt-4.1-mini", Person)
|
|
content = "\n".join(
|
|
[
|
|
openai_result("ok", '{"name": "Ada", "age": 36}'),
|
|
" ",
|
|
openai_result("invalid-model", '{"name": "Ada", "age": "unknown"}'),
|
|
json.dumps({"custom_id": "missing-response"}),
|
|
"not-json",
|
|
]
|
|
)
|
|
|
|
results = processor.parse_results(content)
|
|
|
|
assert len(results) == 4
|
|
assert isinstance(results[0], BatchSuccess)
|
|
assert results[0].custom_id == "ok"
|
|
assert results[0].result == Person(name="Ada", age=36)
|
|
assert isinstance(results[1], BatchError)
|
|
assert results[1].custom_id == "invalid-model"
|
|
assert results[1].error_type == "parsing_error"
|
|
assert "Failed to parse into Person" in results[1].error_message
|
|
assert results[1].raw_data == {"name": "Ada", "age": "unknown"}
|
|
assert isinstance(results[2], BatchError)
|
|
assert results[2].custom_id == "missing-response"
|
|
assert results[2].error_type == "extraction_error"
|
|
assert results[2].error_message == "Unknown error"
|
|
assert isinstance(results[3], BatchError)
|
|
assert results[3].custom_id == "unknown"
|
|
assert results[3].error_type == "json_parse_error"
|
|
assert results[3].raw_data == {"raw_line": "not-json"}
|
|
|
|
|
|
def test_anthropic_results_support_tool_use_and_text_fallback(
|
|
provider: RecordingProvider,
|
|
) -> None:
|
|
del provider
|
|
processor = BatchProcessor("anthropic/claude-sonnet", Person)
|
|
tool_result = {
|
|
"custom_id": "tool",
|
|
"result": {
|
|
"type": "succeeded",
|
|
"message": {
|
|
"content": [
|
|
{"type": "text", "text": "not json"},
|
|
{"type": "tool_use", "input": {"name": "Ada", "age": 36}},
|
|
]
|
|
},
|
|
},
|
|
}
|
|
text_result = {
|
|
"custom_id": "text",
|
|
"result": {
|
|
"type": "succeeded",
|
|
"message": {
|
|
"content": [
|
|
{"type": "text", "text": "not json"},
|
|
{"type": "text", "text": '{"name": "Lin", "age": 28}'},
|
|
]
|
|
},
|
|
},
|
|
}
|
|
|
|
results = processor.parse_results(
|
|
"\n".join([json.dumps(tool_result), json.dumps(text_result)])
|
|
)
|
|
|
|
assert len(results) == 2
|
|
assert isinstance(results[0], BatchSuccess)
|
|
assert isinstance(results[1], BatchSuccess)
|
|
assert results[0].custom_id == "tool"
|
|
assert results[0].result == Person(name="Ada", age=36)
|
|
assert results[1].custom_id == "text"
|
|
assert results[1].result == Person(name="Lin", age=28)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("result", "expected_type", "expected_message"),
|
|
[
|
|
(
|
|
{
|
|
"type": "error",
|
|
"error": {
|
|
"error": {"type": "rate_limit_error", "message": "slow down"}
|
|
},
|
|
},
|
|
"rate_limit_error",
|
|
"slow down",
|
|
),
|
|
(
|
|
{"type": "error", "error": "service unavailable"},
|
|
"anthropic_error",
|
|
"service unavailable",
|
|
),
|
|
(
|
|
{"type": "succeeded", "message": {"content": []}},
|
|
"extraction_error",
|
|
"Unknown error",
|
|
),
|
|
],
|
|
)
|
|
def test_anthropic_error_and_empty_results_keep_provider_details(
|
|
provider: RecordingProvider,
|
|
result: dict[str, Any],
|
|
expected_type: str,
|
|
expected_message: str,
|
|
) -> None:
|
|
del provider
|
|
processor = BatchProcessor("anthropic/claude-sonnet", Person)
|
|
|
|
parsed = processor.parse_results(
|
|
json.dumps({"custom_id": "request-7", "result": result})
|
|
)
|
|
|
|
assert len(parsed) == 1
|
|
assert isinstance(parsed[0], BatchError)
|
|
assert parsed[0].custom_id == "request-7"
|
|
assert parsed[0].error_type == expected_type
|
|
assert parsed[0].error_message == expected_message
|
|
|
|
|
|
def test_extract_returns_none_for_missing_malformed_or_unknown_provider_responses(
|
|
provider: RecordingProvider,
|
|
) -> None:
|
|
del provider
|
|
anthropic = BatchProcessor("anthropic/claude-sonnet", Person)
|
|
openai = BatchProcessor("openai/gpt-4.1-mini", Person)
|
|
unknown = BatchProcessor("local/model", Person)
|
|
|
|
assert anthropic._extract_from_response({"custom_id": "no-result"}) is None
|
|
assert (
|
|
anthropic._extract_from_response(
|
|
{"result": {"type": "succeeded", "message": {"content": "not-a-list"}}}
|
|
)
|
|
is None
|
|
)
|
|
assert anthropic._extract_from_response({"result": {"type": "succeeded"}}) is None
|
|
assert (
|
|
anthropic._extract_from_response(
|
|
{
|
|
"result": {
|
|
"type": "succeeded",
|
|
"message": {"content": [{"type": "image", "source": {}}]},
|
|
}
|
|
}
|
|
)
|
|
is None
|
|
)
|
|
assert anthropic._extract_from_response(
|
|
{
|
|
"result": {
|
|
"type": "succeeded",
|
|
"message": {
|
|
"content": [
|
|
{"type": "image", "source": {}},
|
|
{"type": "text", "text": '{"name":"Ada","age":36}'},
|
|
]
|
|
},
|
|
}
|
|
}
|
|
) == {"name": "Ada", "age": 36}
|
|
assert (
|
|
openai._extract_from_response({"response": {"body": {"choices": []}}}) is None
|
|
)
|
|
assert unknown._extract_from_response({"anything": "goes"}) is None
|