Tests

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

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"

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()

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]}

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