"""Tests for in-memory batch processing functionality.""" import io import json import pytest from pydantic import BaseModel from instructor.batch.request import BatchRequest from instructor.batch.providers.openai import OpenAIProvider from instructor.batch.providers.anthropic import AnthropicProvider # Mark all tests in this module as unit tests (not integration) pytestmark = pytest.mark.unit class User(BaseModel): name: str age: int email: str class TestBatchRequestInMemory: """Test BatchRequest with BytesIO support.""" def test_save_to_bytesio_openai(self): """Test saving BatchRequest to BytesIO for OpenAI format.""" buffer = io.BytesIO() batch_request = BatchRequest[User]( custom_id="test-1", messages=[{"role": "user", "content": "Extract user info"}], response_model=User, model="gpt-4", max_tokens=100, temperature=0.1, ) # Save to BytesIO batch_request.save_to_file(buffer, "openai") # Read back and verify buffer.seek(0) content = buffer.read().decode("utf-8") data = json.loads(content.strip()) assert data["custom_id"] == "test-1" assert data["method"] == "POST" assert data["url"] == "/v1/chat/completions" assert "body" in data assert data["body"]["model"] == "gpt-4" assert "response_format" in data["body"] def test_save_to_bytesio_anthropic(self): """Test saving BatchRequest to BytesIO for Anthropic format.""" buffer = io.BytesIO() batch_request = BatchRequest[User]( custom_id="test-1", messages=[{"role": "user", "content": "Extract user info"}], response_model=User, model="claude-3-sonnet", max_tokens=100, temperature=0.1, ) # Save to BytesIO batch_request.save_to_file(buffer, "anthropic") # Read back and verify buffer.seek(0) content = buffer.read().decode("utf-8") data = json.loads(content.strip()) assert data["custom_id"] == "test-1" assert "params" in data assert data["params"]["model"] == "claude-3-sonnet" assert "tools" in data["params"] def test_save_to_file_still_works(self): """Test that original file-based saving still works.""" import tempfile import os with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".jsonl") as f: temp_path = f.name try: batch_request = BatchRequest[User]( custom_id="test-1", messages=[{"role": "user", "content": "Extract user info"}], response_model=User, model="gpt-4", max_tokens=100, temperature=0.1, ) # Save to file batch_request.save_to_file(temp_path, "openai") # Read back and verify with open(temp_path) as f: content = f.read() data = json.loads(content.strip()) assert data["custom_id"] == "test-1" assert "body" in data finally: if os.path.exists(temp_path): os.unlink(temp_path) def test_multiple_requests_in_buffer(self): """Test writing multiple requests to the same BytesIO buffer.""" buffer = io.BytesIO() for i in range(3): batch_request = BatchRequest[User]( custom_id=f"request-{i}", messages=[{"role": "user", "content": f"Extract user {i}"}], response_model=User, model="gpt-4", max_tokens=100, temperature=0.1, ) batch_request.save_to_file(buffer, "openai") # Read back and verify buffer.seek(0) content = buffer.read().decode("utf-8") lines = [line for line in content.split("\n") if line.strip()] assert len(lines) == 3 for i, line in enumerate(lines): data = json.loads(line) assert data["custom_id"] == f"request-{i}" def test_invalid_buffer_type_raises_error(self): """Test that invalid buffer types raise appropriate errors.""" batch_request = BatchRequest[User]( custom_id="test-1", messages=[{"role": "user", "content": "Extract user info"}], response_model=User, model="gpt-4", max_tokens=100, temperature=0.1, ) with pytest.raises(ValueError, match="Unsupported file_path_or_buffer type"): batch_request.save_to_file( 123, # ty: ignore[invalid-argument-type] - deliberately invalid input "openai", ) class TestProviderInMemorySupport: """Test that providers support BytesIO buffers.""" def test_openai_provider_accepts_bytesio(self): """Test that OpenAI provider accepts BytesIO (without making API calls).""" provider = OpenAIProvider() buffer = io.BytesIO() # Create a valid OpenAI batch request test_data = { "custom_id": "test-1", "method": "POST", "url": "/v1/chat/completions", "body": { "model": "gpt-4", "messages": [{"role": "user", "content": "test"}], "max_tokens": 100, }, } json_line = json.dumps(test_data) + "\n" buffer.write(json_line.encode("utf-8")) buffer.seek(0) # This should not raise a ValueError for unsupported type # (It will raise an exception due to missing API key, but that's expected) try: provider.submit_batch(buffer) except Exception as exc_info: # Make sure it's not a ValueError about unsupported type assert "Unsupported file_path_or_buffer type" not in str(exc_info) def test_anthropic_provider_accepts_bytesio(self): """Test that Anthropic provider accepts BytesIO (without making API calls).""" provider = AnthropicProvider() buffer = io.BytesIO() # Create a valid Anthropic batch request test_data = { "custom_id": "test-1", "params": { "model": "claude-3-sonnet", "messages": [{"role": "user", "content": "test"}], "max_tokens": 100, }, } json_line = json.dumps(test_data) + "\n" buffer.write(json_line.encode("utf-8")) buffer.seek(0) # This should not raise a ValueError for unsupported type # (It will raise an exception due to missing API key, but that's expected) try: provider.submit_batch(buffer) except Exception as exc_info: # Make sure it's not a ValueError about unsupported type assert "Unsupported file_path_or_buffer type" not in str(exc_info) def test_provider_invalid_type_raises_error(self): """Test that providers raise errors for invalid types.""" openai_provider = OpenAIProvider() anthropic_provider = AnthropicProvider() with pytest.raises(ValueError, match="Unsupported file_path_or_buffer type"): openai_provider.submit_batch( 123 # ty: ignore[invalid-argument-type] - deliberately invalid input ) with pytest.raises(ValueError, match="Unsupported file_path_or_buffer type"): anthropic_provider.submit_batch( 123 # ty: ignore[invalid-argument-type] - deliberately invalid input )