Add API retries and concurrent batch processing

tests/test_cli.py

@@ -0,0 +1,54 @@
+"""Tests for the streaming CLI runner."""
+
+from __future__ import annotations
+
+import json
+import threading
+import time
+from pathlib import Path
+
+from swe_data_processing.cli import _run_streaming_stage
+
+
+def test_streaming_stage_processes_records_concurrently(tmp_path: Path) -> None:
+    """Multiple workers run processors in parallel while one thread writes JSONL."""
+
+    input_path = tmp_path / "input.jsonl"
+    output_path = tmp_path / "output.jsonl"
+    input_path.write_text(
+        "".join(
+            json.dumps({"trajectory_id": f"sample-{index}"}) + "\n" for index in range(20)
+        ),
+        encoding="utf-8",
+    )
+    lock = threading.Lock()
+    active = 0
+    peak_active = 0
+
+    def processor(record: dict) -> dict:
+        nonlocal active, peak_active
+        with lock:
+            active += 1
+            peak_active = max(peak_active, active)
+        time.sleep(0.01)
+        with lock:
+            active -= 1
+        return {"sample_id": record["trajectory_id"]}
+
+    status = _run_streaming_stage(
+        input_path=input_path,
+        output_path=output_path,
+        errors_path=None,
+        limit=None,
+        resume=False,
+        workers=4,
+        stage_name="test",
+        processor=processor,
+    )
+
+    written = [json.loads(line) for line in output_path.read_text().splitlines()]
+    assert status == 0
+    assert peak_active == 4
+    assert {value["sample_id"] for value in written} == {
+        f"sample-{index}" for index in range(20)
+    }

tests/test_client.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import json
 
 import httpx
+import pytest
 
 from swe_data_processing.client import GLMClient
 from swe_data_processing.config import Settings
@@ -33,7 +34,7 @@ def test_client_parses_fenced_json_and_never_sends_key_in_body() -> None:
         "required": ["status"],
         "properties": {"status": {"const": "ok"}},
     }
-    settings = Settings(api_key="test-secret", max_retries=1)
+    settings = Settings(api_key="test-secret", max_retries=0)
     with GLMClient(settings, transport=httpx.MockTransport(handler)) as client:
         response = client.invoke_json(system_prompt="Return JSON.", payload={"ping": True}, schema=schema)
 
@@ -62,10 +63,41 @@ def test_client_falls_back_when_vendor_extensions_are_rejected() -> None:
         "required": ["status"],
         "properties": {"status": {"const": "ok"}},
     }
-    settings = Settings(api_key="test-secret", max_retries=2)
+    settings = Settings(api_key="test-secret", max_retries=1)
     with GLMClient(settings, transport=httpx.MockTransport(handler)) as client:
         response = client.invoke_json(system_prompt="Return JSON.", payload={}, schema=schema)
 
     assert "thinking" in calls[0]
     assert "thinking" not in calls[1]
     assert response.compatibility_fallback_used is True
+
+
+def test_client_retries_five_failures_before_succeeding(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Five retryable failures are followed by a sixth and final attempt."""
+
+    calls = 0
+
+    def handler(_: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        if calls <= 5:
+            return httpx.Response(503, json={"error": {"message": "temporarily unavailable"}})
+        return httpx.Response(
+            200,
+            json={"choices": [{"message": {"content": "{\"status\":\"ok\"}"}}]},
+        )
+
+    monkeypatch.setattr("swe_data_processing.client.time.sleep", lambda _: None)
+    schema = {
+        "type": "object",
+        "required": ["status"],
+        "properties": {"status": {"const": "ok"}},
+    }
+    settings = Settings(api_key="test-secret", max_retries=5)
+    with GLMClient(settings, transport=httpx.MockTransport(handler)) as client:
+        response = client.invoke_json(system_prompt="Return JSON.", payload={}, schema=schema)
+
+    assert calls == 6
+    assert response.data == {"status": "ok"}

tests/test_config.py

@@ -24,4 +24,6 @@ def test_missing_api_key_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
 
 def test_feature_only_settings_can_omit_key(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.delenv("GLM_API_KEY", raising=False)
-    assert Settings.from_env(require_api_key=False).api_key == ""
+    settings = Settings.from_env(require_api_key=False)
+    assert settings.api_key == ""
+    assert settings.max_retries == 5