200 lines
6.1 KiB
Python
200 lines
6.1 KiB
Python
import asyncio
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from great_ai import GreatAI
|
|
from great_ai.context import get_context
|
|
|
|
|
|
def test_nested_tuple_inputs_can_be_cached(isolated_context: None) -> None:
|
|
calls = 0
|
|
|
|
@GreatAI.create
|
|
def predict(value: Tuple[Dict[str, int]]) -> int:
|
|
nonlocal calls
|
|
calls += 1
|
|
return value[0]["x"]
|
|
|
|
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
|
|
|
|
|
|
def test_prediction_mutation_does_not_change_cache_keys(isolated_context: None) -> None:
|
|
calls = 0
|
|
|
|
@GreatAI.create
|
|
def predict(values: List[int]) -> int:
|
|
nonlocal calls
|
|
calls += 1
|
|
values.append(1)
|
|
return len(values)
|
|
|
|
original: List[int] = []
|
|
assert predict(original).output == 1
|
|
assert predict([]).output == 1
|
|
assert original == []
|
|
assert calls == 1
|
|
|
|
|
|
def test_caller_mutation_does_not_change_cached_trace_values(
|
|
isolated_context: None,
|
|
) -> None:
|
|
@GreatAI.create
|
|
def predict(values: List[List[int]]) -> int:
|
|
return values[0][0]
|
|
|
|
original = [[1]]
|
|
first = predict(original)
|
|
original[0][0] = 2
|
|
assert predict(original).output == 2
|
|
assert predict([[1]]).trace_id == first.trace_id
|
|
assert first.logged_values["arg:values:value"] == [[1]]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("is_async", [False, True])
|
|
@pytest.mark.parametrize("do_not_persist", [False, True])
|
|
async def test_returned_traces_cannot_mutate_cached_results(
|
|
isolated_context: None, is_async: bool, do_not_persist: bool
|
|
) -> None:
|
|
calls = 0
|
|
|
|
def predict(values: List[int]) -> List[List[int]]:
|
|
nonlocal calls
|
|
calls += 1
|
|
return [values]
|
|
|
|
async def predict_async(values: List[int]) -> List[List[int]]:
|
|
return predict(values)
|
|
|
|
service: Any = GreatAI.create(predict_async if is_async else predict)
|
|
first = service([1], do_not_persist_traces=do_not_persist)
|
|
if is_async:
|
|
first = await first
|
|
trace_id = first.trace_id
|
|
first.output[0].append(99)
|
|
first.logged_values["arg:values:value"].append(88)
|
|
first.tags.append("changed")
|
|
first.trace_id = "changed"
|
|
first.feedback = "changed"
|
|
|
|
second = service([1], do_not_persist_traces=do_not_persist)
|
|
if is_async:
|
|
second = await second
|
|
assert second.trace_id == trace_id
|
|
assert second.output == [[1]]
|
|
assert second.logged_values["arg:values:value"] == [1]
|
|
assert "changed" not in second.tags
|
|
assert second.feedback is None
|
|
assert calls == 1
|
|
stored = get_context().tracing_database.get(trace_id)
|
|
if do_not_persist:
|
|
assert stored is None
|
|
else:
|
|
assert stored is not None
|
|
assert stored.output == [[1]]
|
|
|
|
|
|
@pytest.mark.parametrize("is_async", [False, True])
|
|
def test_cached_predictions_follow_http_feedback_and_deletion(
|
|
isolated_context: None, is_async: bool
|
|
) -> None:
|
|
calls = 0
|
|
|
|
def predict(value: int) -> int:
|
|
nonlocal calls
|
|
calls += 1
|
|
return value * 2
|
|
|
|
async def predict_async(value: int) -> int:
|
|
return predict(value)
|
|
|
|
service: Any = GreatAI.create(predict_async if is_async else predict)
|
|
with TestClient(service.app) as client:
|
|
first = client.post("/predict", json={"value": 2}).json()
|
|
trace_url = f"/traces/{first['trace_id']}"
|
|
feedback_url = f"{trace_url}/feedback/"
|
|
assert (
|
|
client.put(
|
|
feedback_url, json={"feedback": {"label": "correct"}}
|
|
).status_code
|
|
== 202
|
|
)
|
|
second = client.post("/predict", json={"value": 2}).json()
|
|
assert second["trace_id"] == first["trace_id"]
|
|
assert second["feedback"] == {"label": "correct"}
|
|
assert second == client.get(trace_url).json()
|
|
assert calls == 1
|
|
assert client.delete(feedback_url).status_code == 204
|
|
assert client.post("/predict", json={"value": 2}).json()["feedback"] is None
|
|
assert client.delete(trace_url).status_code == 204
|
|
third = client.post("/predict", json={"value": 2}).json()
|
|
assert third["trace_id"] != first["trace_id"]
|
|
assert third["output"] == 4
|
|
assert calls == 2
|
|
assert client.get(f"/traces/{third['trace_id']}").json() == third
|
|
assert (
|
|
client.put(
|
|
f"/traces/{third['trace_id']}/feedback/", json={"feedback": "correct"}
|
|
).status_code
|
|
== 202
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("is_async", [False, True])
|
|
async def test_cached_predictions_observe_direct_database_changes(
|
|
isolated_context: None, is_async: bool
|
|
) -> None:
|
|
def predict(value: int) -> int:
|
|
return value
|
|
|
|
async def predict_async(value: int) -> int:
|
|
return value
|
|
|
|
service: Any = GreatAI.create(predict_async if is_async else predict)
|
|
first = service(1)
|
|
if is_async:
|
|
first = await first
|
|
database = get_context().tracing_database
|
|
stored = database.get(first.trace_id)
|
|
assert stored is not None
|
|
stored.feedback = ["updated outside the service"]
|
|
database.update(first.trace_id, stored)
|
|
second = service(1)
|
|
if is_async:
|
|
second = await second
|
|
assert second.feedback == stored.feedback
|
|
second.feedback.append("caller change")
|
|
database.delete(first.trace_id)
|
|
third = service(1)
|
|
if is_async:
|
|
third = await third
|
|
assert third.trace_id != first.trace_id
|
|
assert database.get(third.trace_id) is not None
|