Retry failed and cancelled async predictions instead of caching errors

tests/test_cache_regressions.py

@@ -1,5 +1,8 @@
+import asyncio
 from typing import Dict, Tuple
 
+import pytest
+
 from great_ai import GreatAI
 
 
@@ -15,3 +18,27 @@ def test_nested_tuple_inputs_can_be_cached(isolated_context: None) -> None:
     assert predict(({"x": 7},)).output == 7
     assert predict(({"x": 7},)).output == 7
     assert calls == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cancelled", [False, True])
+async def test_async_predictions_retry_after_failure(
+    isolated_context: None, cancelled: bool
+) -> None:
+    calls = 0
+
+    @GreatAI.create
+    async def predict(value: int) -> int:
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            if cancelled:
+                raise asyncio.CancelledError()
+            raise ValueError("temporary failure")
+        return value
+
+    with pytest.raises(asyncio.CancelledError if cancelled else ValueError):
+        await predict(7)
+    assert (await predict(7)).output == 7
+    assert (await predict(7)).output == 7
+    assert calls == 2