diff --git a/tests/test_cache_regressions.py b/tests/test_cache_regressions.py index a84573e..70e20ba 100644 --- a/tests/test_cache_regressions.py +++ b/tests/test_cache_regressions.py @@ -1,8 +1,10 @@ import asyncio -from typing import Dict, List, Tuple +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: @@ -73,3 +75,126 @@ def test_caller_mutation_does_not_change_cached_trace_values( 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 diff --git a/tests/test_feedback_training_inputs.py b/tests/test_feedback_training_inputs.py new file mode 100644 index 0000000..497881f --- /dev/null +++ b/tests/test_feedback_training_inputs.py @@ -0,0 +1,89 @@ +from typing import Any, Dict + +from fastapi.testclient import TestClient +from great_ai import ( + GreatAI, + Trace, + add_ground_truth, + log_metric, + parameter, + query_ground_truth, + save_model, + use_model, +) + + +def test_uploaded_and_online_feedback_share_training_input_shape( + isolated_context: None, +) -> None: + @GreatAI.create + def predict(text: str) -> str: + log_metric("length", len(text)) + return text.upper() + + add_ground_truth(["dog"], ["DOG"]) + with TestClient(predict.app) as client: + trace = client.post("/predict", json={"text": "cat"}).json() + client.put(f"/traces/{trace['trace_id']}/feedback/", json={"feedback": "CAT"}) + pairs = {(trace.input, trace.feedback) for trace in query_ground_truth()} + assert pairs == {("cat", "CAT"), ("dog", "DOG")} + + +def test_dictionary_inputs_remain_dictionaries(isolated_context: None) -> None: + @GreatAI.create + def predict(data: Dict[str, Any]) -> int: + return len(data) + + original = {"input": {"nested": [1, 2]}} + add_ground_truth([original], [1], tags="online") + with TestClient(predict.app) as client: + trace = client.post("/predict", json={"data": original}).json() + client.put(f"/traces/{trace['trace_id']}/feedback/", json={"feedback": 1}) + data = query_ground_truth() + assert len(data) == 2 + assert all(trace.input == original for trace in data) + + +def test_multi_argument_inputs_exclude_models_metrics_and_private_arguments( + isolated_context: None, +) -> None: + save_model(3, "constant") + + @GreatAI.create + @parameter("secret", disable_logging=True) + @use_model("constant") + def predict(text: str, model: int, *, secret: str, count: int = 2) -> str: + log_metric("custom", model) + return text * count + + original = {"text": "cat", "count": 2} + add_ground_truth([original], ["catcat"]) + with TestClient(predict.app) as client: + trace = client.post( + "/predict", json={"text": "cat", "secret": "private"} + ).json() + client.put( + f"/traces/{trace['trace_id']}/feedback/", json={"feedback": "catcat"} + ) + assert all(trace.input == original for trace in query_ground_truth()) + assert "private" not in str(query_ground_truth()) + + +def test_legacy_uploaded_dictionary_and_online_traces_remain_readable() -> None: + base: Dict[str, Any] = dict( + trace_id="old", created="2022-01-01", original_execution_time_ms=0, models=[] + ) + uploaded = Trace( + **base, logged_values={"text": "cat", "count": 2}, tags=["ground_truth"] + ) + assert uploaded.input == {"text": "cat", "count": 2} + online = Trace( + **base, + logged_values={ + "arg:text:value": "cat", + "arg:text:length": 3, + "metric:predict:custom": 1, + }, + tags=["online"], + ) + assert online.input == "cat" diff --git a/tests/test_notebook_setup.py b/tests/test_notebook_setup.py new file mode 100644 index 0000000..2c6af5f --- /dev/null +++ b/tests/test_notebook_setup.py @@ -0,0 +1,42 @@ +import json +from pathlib import Path +from typing import Any, Dict + +import pytest +from great_ai import ParallelTinyDbDriver +from great_ai.context import get_context +from great_ai.large_file import LargeFileLocal + +ROOT = Path(__file__).resolve().parents[1] +NOTEBOOKS = sorted((ROOT / "docs").rglob("*.ipynb")) + + +@pytest.mark.parametrize( + "notebook", NOTEBOOKS, ids=lambda path: str(path.relative_to(ROOT)) +) +def test_notebook_setup_uses_checkout_paths_from_another_directory( + isolated_context: None, + notebook: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.syspath_prepend(str(ROOT)) + from docs import notebook_support + + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("GREAT_AI_NOTEBOOK_WORKDIR", str(tmp_path)) + monkeypatch.setattr(notebook_support, "FULL_RUN", False) + monkeypatch.setattr(ParallelTinyDbDriver, "path_to_db", tmp_path / "traces.json") + cells = json.loads(notebook.read_text())["cells"] + setup = next( + "".join(cell["source"]) for cell in cells if cell["cell_type"] == "code" + ) + namespace: Dict[str, Any] = {} + exec(compile(setup, str(notebook), "exec"), namespace) + work_dir = tmp_path / notebook.parent.name + assert work_dir.is_dir() + assert LargeFileLocal.cache_path == work_dir / "models-cache" + database = get_context().tracing_database + assert isinstance(database, ParallelTinyDbDriver) + assert database.path_to_db == work_dir / "tracing_database.json" + assert (notebook_support.DATA_DIR / "train.txt").is_file() diff --git a/tests/test_numpy_predictions.py b/tests/test_numpy_predictions.py new file mode 100644 index 0000000..e470d02 --- /dev/null +++ b/tests/test_numpy_predictions.py @@ -0,0 +1,194 @@ +from typing import Any, Dict, Optional +from unittest.mock import patch + +import numpy as np +import pytest +from fastapi.testclient import TestClient +from great_ai import GreatAI, log_metric +from great_ai.context import get_context +from great_ai.persistence.mongodb_driver import MongoDbDriver +from numpy.typing import NDArray +from pydantic import BaseModel, ConfigDict, Field + + +class ArraySubclass(np.ndarray): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("cache_size", [0, 8]) +async def test_numpy_arrays_work_with_and_without_caching( + isolated_context: None, is_async: bool, cache_size: int +) -> None: + get_context().prediction_cache_size = cache_size + calls = 0 + + def predict(values: np.ndarray) -> np.ndarray: + nonlocal calls + calls += 1 + log_metric("mean", np.float32(values.mean())) + return values * 2 + + async def predict_async(values: np.ndarray) -> np.ndarray: + return predict(values) + + service: Any = GreatAI.create(predict_async if is_async else predict) + original = np.array([1, 2], dtype=np.int64) + first = service(original) + if is_async: + first = await first + original[0] = 9 + first.output[0] = 99 + second = service(np.array([1, 2], dtype=np.int64)) + if is_async: + second = await second + np.testing.assert_array_equal(second.output, [2, 4]) + assert second.output.dtype == np.int64 + assert calls == (1 if cache_size else 2) + stored = get_context().tracing_database.get(second.trace_id) + assert stored is not None + assert stored.output == [2, 4] + assert stored.input == [1, 2] + assert stored.logged_values["metric:predict:mean"] == 1.5 + + +def test_numpy_cache_keys_preserve_shape_dtype_and_object_values( + isolated_context: None, +) -> None: + calls = 0 + + @GreatAI.create + def predict(values: np.ndarray) -> str: + nonlocal calls + calls += 1 + return str(values.shape) + str(values.dtype) + + for values in ( + np.array([1, 2], dtype=np.int32), + np.array([1, 2], dtype=np.int64), + np.array([1, 2], dtype=np.int64).view(ArraySubclass), + np.array([[1, 2]], dtype=np.int64), + np.array([{"x": [1]}], dtype=object), + ): + first = predict(values) + assert predict(values.copy()).trace_id == first.trace_id + assert calls == 5 + + +def test_http_numpy_arrays_have_schemas_and_round_trip_through_storage( + isolated_context: None, +) -> None: + @GreatAI.create + def predict(values: Optional[NDArray[np.float64]]) -> list[np.ndarray]: + assert isinstance(values, np.ndarray) + log_metric( + "nested", + {"maximum": np.int64(values.max()), "flags": np.array([True, False])}, + ) + return [values * 2] + + with TestClient(predict.app) as client: + response = client.post("/predict", json={"values": [1, 2]}) + assert response.status_code == 200 + trace = response.json() + assert trace["output"] == [[2, 4]] + assert trace["logged_values"]["metric:predict:nested"] == { + "maximum": 2, + "flags": [True, False], + } + trace_url = f"/traces/{trace['trace_id']}" + assert client.get(trace_url).json() == trace + assert client.post("/traces", json={}).json() == [trace] + assert client.post("/predict", json={"values": "wrong"}).status_code == 422 + schema = client.get("/openapi.json").json() + assert ( + schema["components"]["schemas"]["InputModel"]["properties"]["values"][ + "anyOf" + ][0]["type"] + == "array" + ) + response_ref = schema["paths"]["/predict"]["post"]["responses"]["200"][ + "content" + ]["application/json"]["schema"]["$ref"].split("/")[-1] + output = schema["components"]["schemas"][response_ref]["properties"]["output"][ + "anyOf" + ][0] + assert output["type"] == output["items"]["type"] == "array" + + +def test_numpy_scalar_annotations_are_preserved_in_python_and_openapi( + isolated_context: None, +) -> None: + @GreatAI.create + def predict(value: np.int64) -> np.float32: + assert isinstance(value, np.int64) + return np.float32(value / 2) + + assert isinstance(predict(np.int64(3)).output, np.float32) + with TestClient(predict.app) as client: + response = client.post("/predict", json={"value": 3}) + assert response.status_code == 200 + assert response.json()["output"] == 1.5 + schema = client.get("/openapi.json").json() + assert ( + schema["components"]["schemas"]["InputModel"]["properties"]["value"]["type"] + == "integer" + ) + response_ref = schema["paths"]["/predict"]["post"]["responses"]["200"][ + "content" + ]["application/json"]["schema"]["$ref"].split("/")[-1] + assert ( + schema["components"]["schemas"][response_ref]["properties"]["output"][ + "anyOf" + ][0]["type"] + == "number" + ) + + +def test_mongodb_serializes_numpy_values_for_insert_batch_and_update( + isolated_context: None, +) -> None: + @GreatAI.create + def predict(values: np.ndarray) -> Dict[str, Any]: + return {"values": values, "count": np.int64(values.size)} + + trace = predict(np.array([1, 2])) + driver = MongoDbDriver.__new__(MongoDbDriver) + driver.mongo_connection_string = "mongodb://127.0.0.1:1" + driver.mongo_database = "test" + with patch("great_ai.persistence.mongodb_driver.MongoClient") as factory: + collection = ( + factory.__getitem__.return_value.return_value.__enter__.return_value[ + "test" + ].traces + ) + driver.save(trace) + driver.save_batch([trace]) + driver.update(trace.trace_id, trace) + inserted = collection.insert_one.call_args.args[0] + assert inserted["output"] == {"values": [1, 2], "count": 2} + assert inserted["logged_values"]["arg:values:value"] == [1, 2] + assert collection.insert_many.call_args.args[0] == [inserted] + assert collection.replace_one.call_args.args[1] == inserted + + +def test_numpy_serialization_preserves_exclusions_and_model_aliases( + isolated_context: None, +) -> None: + class Result(BaseModel): + values: Any = Field(alias="array") + secret: str = "private" + model_config = ConfigDict(populate_by_name=True) + + @GreatAI.create + def predict(value: int) -> Result: + return Result(array=np.array([value]), secret="private") + + trace = predict(1) + dumped = trace.model_dump(mode="json", exclude={"output": {"secret"}}) + assert dumped["output"] == {"values": [1]} + aliased = trace.model_dump( + mode="json", by_alias=True, exclude={"output": {"secret"}} + ) + assert aliased["output"] == {"array": [1]} diff --git a/tests/test_prediction_contract.py b/tests/test_prediction_contract.py new file mode 100644 index 0000000..8082b1f --- /dev/null +++ b/tests/test_prediction_contract.py @@ -0,0 +1,80 @@ +from typing import List + +import pytest +from fastapi.testclient import TestClient +from great_ai import GreatAI, parameter +from pydantic import BaseModel + + +class Score(BaseModel): + label: str + confidence: float + + +@pytest.mark.parametrize("is_async", [False, True]) +def test_custom_validation_returns_client_error_but_prediction_failures_do_not( + isolated_context: None, is_async: bool +) -> None: + @parameter("value", validate=lambda value: value > 0) + def predict(value: int) -> int: + if value == 7: + raise ValueError("prediction failed") + return value + + @parameter("value", validate=lambda value: value > 0) + async def predict_async(value: int) -> int: + return predict(value) + + service = GreatAI.create(predict_async if is_async else predict) + with TestClient(service.app) as client: + rejected = client.post("/predict", json={"value": -1}) + assert rejected.status_code == 422 + assert "did not pass validation" in rejected.json()["detail"] + assert client.post("/predict", json={"value": "wrong"}).status_code == 422 + assert client.post("/predict", json={"value": 7}).status_code == 500 + valid = client.post("/predict", json={"value": 1}) + assert valid.status_code == 200 + assert valid.json()["output"] == 1 + + +@pytest.mark.parametrize("is_async", [False, True]) +def test_openapi_preserves_typed_nested_prediction_outputs( + isolated_context: None, is_async: bool +) -> None: + def predict(text: str) -> List[Score]: + return [Score(label=text, confidence=0.75)] + + async def predict_async(text: str) -> List[Score]: + return predict(text) + + service = GreatAI.create(predict_async if is_async else predict) + with TestClient(service.app) as client: + schema = client.get("/openapi.json").json() + response_ref = schema["paths"]["/predict"]["post"]["responses"]["200"][ + "content" + ]["application/json"]["schema"]["$ref"].split("/")[-1] + output = schema["components"]["schemas"][response_ref]["properties"]["output"][ + "anyOf" + ][0] + assert output == { + "items": {"$ref": "#/components/schemas/Score"}, + "type": "array", + } + assert ( + schema["components"]["schemas"]["Score"]["properties"]["confidence"]["type"] + == "number" + ) + response = client.post("/predict", json={"text": "cat"}) + assert response.status_code == 200 + assert response.json()["output"] == [{"label": "cat", "confidence": 0.75}] + + +def test_http_rejects_outputs_that_violate_the_return_annotation( + isolated_context: None, +) -> None: + @GreatAI.create + def predict(value: int) -> int: + return "invalid" # type: ignore[return-value] + + with TestClient(predict.app, raise_server_exceptions=False) as client: + assert client.post("/predict", json={"value": 1}).status_code == 500