Move file content to be a reference in state file (#27)

* Move file content to be a reference in state file

* Clear pending when done in state file

* Fix tests locally

* Fix warning -> debug

Author: agamm

batchata/core/batch.py

@@ -1,5 +1,6 @@
 """Batch builder."""
 
+import logging
 import uuid
 from datetime import datetime
 from pathlib import Path
@@ -281,6 +282,14 @@ def add_job(
         if isinstance(file, str):
             file = Path(file)
 
+        # Warn about temporary file paths that may not persist
+        if file:
+            file_str = str(file)
+            if "/tmp/" in file_str or "/var/folders/" in file_str or "temp" in file_str.lower():
+                logger = logging.getLogger("batchata")
+                logger.debug(f"File path appears to be in a temporary directory: {file}")
+                logger.debug("This may cause issues when resuming from state if temp files are cleaned up")
+        
         # Create job
         job = Job(
             id=job_id,

batchata/core/batch_run.py

@@ -132,12 +132,27 @@ def _resume_from_state(self):
         self.pending_jobs = []
         for job_data in state.pending_jobs:
             job = Job.from_dict(job_data)
-            self.pending_jobs.append(job)
+            # Check if file exists (if job has a file)
+            if job.file and not job.file.exists():
+                logger.error(f"File not found for job {job.id}: {job.file}")
+                logger.error("This may happen if files were in temporary directories that were cleaned up")
+                self.failed_jobs[job.id] = f"File not found: {job.file}"
+            else:
+                self.pending_jobs.append(job)
 
-        # Restore completed results
-        for result_data in state.completed_results:
-            result = JobResult.from_dict(result_data)
-            self.completed_results[result.job_id] = result
+        # Restore completed results from file references
+        for result_ref in state.completed_results:
+            job_id = result_ref["job_id"]
+            file_path = result_ref["file_path"]
+            try:
+                with open(file_path, 'r') as f:
+                    result_data = json.load(f)
+                result = JobResult.from_dict(result_data)
+                self.completed_results[job_id] = result
+            except Exception as e:
+                logger.error(f"Failed to load result for {job_id} from {file_path}: {e}")
+                # Move to failed jobs if we can't load the result
+                self.failed_jobs[job_id] = f"Failed to load result file: {e}"
 
         # Restore failed jobs
         for job_data in state.failed_jobs:
@@ -162,7 +177,10 @@ def to_json(self) -> Dict:
         return {
             "created_at": datetime.now().isoformat(),
             "pending_jobs": [job.to_dict() for job in self.pending_jobs],
-            "completed_results": [result.to_dict() for result in self.completed_results.values()],
+            "completed_results": [
+                {"job_id": job_id, "file_path": str(self.results_dir / f"{job_id}.json")}
+                for job_id in self.completed_results.keys()
+            ],
             "failed_jobs": [
                 {
                     "id": job_id, 
@@ -505,10 +523,15 @@ def _update_batch_results(self, batch_result: Dict):
                 self.failed_jobs[result.job_id] = error_message
                 self._save_result_to_file(result)
                 logger.error(f"✗ Job {result.job_id} failed: {result.error}")
+            
+            # Remove completed/failed job from pending
+            self.pending_jobs = [job for job in self.pending_jobs if job.id != result.job_id]
 
         # Update failed jobs
         for job_id, error in failed.items():
             self.failed_jobs[job_id] = error
+            # Remove failed job from pending
+            self.pending_jobs = [job for job in self.pending_jobs if job.id != job_id]
             logger.error(f"✗ Job {job_id} failed: {error}")
 
         # Update batch tracking

batchata/utils/state.py

@@ -22,7 +22,7 @@ class BatchState:
         created_at: When the batch was created
         pending_jobs: Jobs that haven't been submitted yet
         active_batches: Currently running batch IDs
-        completed_results: Results from completed jobs
+        completed_results: File references to completed job results
         failed_jobs: Jobs that failed with errors
         total_cost_usd: Total cost incurred so far
         config: Original batch configuration
@@ -31,7 +31,7 @@ class BatchState:
     created_at: str  # ISO format datetime
     pending_jobs: List[Dict[str, Any]]  # Serialized Job objects
     active_batches: List[str]  # Provider batch IDs
-    completed_results: List[Dict[str, Any]]  # Serialized JobResult objects
+    completed_results: List[Dict[str, str]]  # File references: [{"job_id": "job_123", "file_path": "/path/to/result.json"}]
     failed_jobs: List[Dict[str, Any]]  # Jobs with error info
     total_cost_usd: float
     config: Dict[str, Any]  # Batch configuration

tests/core/test_batch_validation.py

@@ -1,8 +1,10 @@
 """Tests for batch validation with citations."""
 
 import pytest
+import os
 from pydantic import BaseModel
 from typing import Optional
+from unittest.mock import patch
 
 from batchata import Batch
 
@@ -29,6 +31,15 @@ class NestedInvoice(BaseModel):
 class TestBatchCitationValidation:
     """Test early validation of citation compatibility."""
 
+    @pytest.fixture(autouse=True)
+    def mock_api_keys(self):
+        """Provide mock API keys for provider initialization."""
+        with patch.dict(os.environ, {
+            'ANTHROPIC_API_KEY': 'test-anthropic-key',
+            'OPENAI_API_KEY': 'test-openai-key'
+        }):
+            yield
+    
     def test_flat_model_with_citations_allowed(self):
         """Test that flat models work with citations."""
         batch = Batch("./results").set_state(file="./state")

tests/core/test_pdf_validation.py

@@ -1,8 +1,10 @@
 """Tests for PDF validation in batch jobs."""
 
 import pytest
+import os
 import tempfile
 from pathlib import Path
+from unittest.mock import patch
 
 from batchata.core.batch import Batch
 from batchata.utils.pdf import create_pdf
@@ -12,14 +14,23 @@
 class TestPdfValidation:
     """Test PDF validation when citations are enabled."""
 
+    @pytest.fixture(autouse=True)
+    def mock_api_keys(self):
+        """Provide mock API keys for provider initialization."""
+        with patch.dict(os.environ, {
+            'ANTHROPIC_API_KEY': 'test-anthropic-key',
+            'OPENAI_API_KEY': 'test-openai-key'
+        }):
+            yield
+    
     def test_image_pdf_with_citations_fails(self):
         """Test image PDF with citations should fail."""
         with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as tmp:
             tmp.write(b"not a real pdf")  # Score 0.0
             tmp.flush()
 
             batch = Batch("/tmp/results").set_state(file="/tmp/state.json")
-            batch.set_default_params(model="claude-3-5-sonnet-latest")
+            batch.set_default_params(model="claude-3-5-sonnet-20241022")
 
             with pytest.raises(ValidationError, match="appears to be image-only"):
                 batch.add_job(file=tmp.name, prompt="Test", enable_citations=True)
@@ -33,7 +44,7 @@ def test_image_pdf_without_citations_works(self):
             tmp.flush()
 
             batch = Batch("/tmp/results").set_state(file="/tmp/state.json")
-            batch.set_default_params(model="claude-3-5-sonnet-latest")
+            batch.set_default_params(model="claude-3-5-sonnet-20241022")
             batch.add_job(file=tmp.name, prompt="Test", enable_citations=False)
 
             assert len(batch.jobs) == 1
@@ -49,7 +60,7 @@ def test_textual_pdf_with_citations_works(self):
             tmp.flush()
 
             batch = Batch("/tmp/results").set_state(file="/tmp/state.json")
-            batch.set_default_params(model="claude-3-5-sonnet-latest")
+            batch.set_default_params(model="claude-3-5-sonnet-20241022")
             batch.add_job(file=tmp.name, prompt="Test", enable_citations=True)
 
             assert len(batch.jobs) == 1
@@ -62,7 +73,7 @@ def test_non_pdf_bypasses_validation(self):
             tmp.flush()
 
             batch = Batch("/tmp/results").set_state(file="/tmp/state.json")
-            batch.set_default_params(model="claude-3-5-sonnet-latest")
+            batch.set_default_params(model="claude-3-5-sonnet-20241022")
             batch.add_job(file=tmp.name, prompt="Test", enable_citations=True)
 
             assert len(batch.jobs) == 1

tests/providers/test_provider_registry.py

@@ -7,6 +7,7 @@
 """
 
 import pytest
+import os
 from unittest.mock import patch, MagicMock
 
 from batchata.providers import get_provider
@@ -18,6 +19,15 @@
 class TestProviderRegistry:
     """Test provider registry functionality."""
 
+    @pytest.fixture(autouse=True)
+    def mock_api_keys(self):
+        """Provide mock API keys for provider initialization."""
+        with patch.dict(os.environ, {
+            'ANTHROPIC_API_KEY': 'test-anthropic-key',
+            'OPENAI_API_KEY': 'test-openai-key'
+        }):
+            yield
+    
     def test_provider_lookup_by_model(self):
         """Test looking up providers by model name."""
         # Create a mock provider and register it

tests/utils/test_state.py

@@ -35,8 +35,7 @@ def test_save_and_load_state(self, temp_dir):
             ],
             active_batches=["batch-123"],
             completed_results=[
-                {"job_id": "job-0", "content": "A0", "cost": 0.01, 
-                 "input_tokens": 10, "output_tokens": 20}
+                {"job_id": "job-0", "file_path": "/tmp/job-0.json"}
             ],
             failed_jobs=[],
             total_cost_usd=0.01,
@@ -57,6 +56,7 @@ def test_save_and_load_state(self, temp_dir):
         assert loaded.active_batches[0] == "batch-123"
         assert len(loaded.completed_results) == 1
         assert loaded.completed_results[0]["job_id"] == "job-0"
+        assert loaded.completed_results[0]["file_path"] == "/tmp/job-0.json"
         assert loaded.total_cost_usd == 0.01
         assert loaded.config["max_parallel_batches"] == 10