orca-sdk 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

orca_sdk/memoryset_test.py

@@ -5,10 +5,11 @@ import pytest
 from datasets.arrow_dataset import Dataset
 
 from .classification_model import ClassificationModel
-from .conftest import skip_in_prod
+from .conftest import skip_in_ci, skip_in_prod
 from .datasource import Datasource
 from .embedding_model import PretrainedEmbeddingModel
 from .memoryset import LabeledMemoryset, ScoredMemory, ScoredMemoryset, Status
+from .regression_model import RegressionModel
 
 """
 Test Performance Note:
@@ -39,9 +40,10 @@ def test_create_memoryset(readonly_memoryset: LabeledMemoryset, hf_dataset: Data
     assert readonly_memoryset.index_params == {"n_lists": 100}
 
 
-def test_create_memoryset_unauthenticated(unauthenticated, datasource):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        LabeledMemoryset.create("test_memoryset", datasource)
+def test_create_memoryset_unauthenticated(unauthenticated_client, datasource):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            LabeledMemoryset.create("test_memoryset", datasource)
 
 
 def test_create_memoryset_invalid_input(datasource):
@@ -87,6 +89,75 @@ def test_create_memoryset_already_exists_open(hf_dataset, label_names, readonly_
     assert opened_memoryset.length == len(hf_dataset)
 
 
+def test_if_exists_error_no_datasource_creation(
+    readonly_memoryset: LabeledMemoryset,
+):
+    memoryset_name = readonly_memoryset.name
+    datasource_name = f"{memoryset_name}_datasource"
+    Datasource.drop(datasource_name, if_not_exists="ignore")
+    assert not Datasource.exists(datasource_name)
+    with pytest.raises(ValueError):
+        LabeledMemoryset.from_list(memoryset_name, [{"value": "new value", "label": 0}], if_exists="error")
+    assert not Datasource.exists(datasource_name)
+
+
+def test_if_exists_open_reuses_existing_datasource(
+    readonly_memoryset: LabeledMemoryset,
+):
+    memoryset_name = readonly_memoryset.name
+    datasource_name = f"{memoryset_name}_datasource"
+    Datasource.drop(datasource_name, if_not_exists="ignore")
+    assert not Datasource.exists(datasource_name)
+    reopened = LabeledMemoryset.from_list(memoryset_name, [{"value": "new value", "label": 0}], if_exists="open")
+    assert reopened.id == readonly_memoryset.id
+    assert not Datasource.exists(datasource_name)
+
+
+def test_create_memoryset_string_label():
+    assert not LabeledMemoryset.exists("test_string_label")
+    memoryset = LabeledMemoryset.from_hf_dataset(
+        "test_string_label",
+        Dataset.from_dict({"value": ["terrible", "great"], "label": ["negative", "positive"]}),
+    )
+    assert memoryset is not None
+    assert memoryset.length == 2
+    assert memoryset.label_names == ["negative", "positive"]
+    assert memoryset[0].label == 0
+    assert memoryset[1].label == 1
+    assert memoryset[0].label_name == "negative"
+    assert memoryset[1].label_name == "positive"
+
+
+def test_create_memoryset_integer_label():
+    assert not LabeledMemoryset.exists("test_integer_label")
+    memoryset = LabeledMemoryset.from_hf_dataset(
+        "test_integer_label",
+        Dataset.from_dict({"value": ["terrible", "great"], "label": [0, 1]}),
+        label_names=["negative", "positive"],
+    )
+    assert memoryset is not None
+    assert memoryset.length == 2
+    assert memoryset.label_names == ["negative", "positive"]
+    assert memoryset[0].label == 0
+    assert memoryset[1].label == 1
+    assert memoryset[0].label_name == "negative"
+    assert memoryset[1].label_name == "positive"
+
+
+def test_create_memoryset_null_labels():
+    memoryset = LabeledMemoryset.from_hf_dataset(
+        "test_null_labels",
+        Dataset.from_dict({"value": ["terrible", "great"]}),
+        label_names=["negative", "positive"],
+        label_column=None,
+    )
+    assert memoryset is not None
+    assert memoryset.length == 2
+    assert memoryset.label_names == ["negative", "positive"]
+    assert memoryset[0].label == None
+    assert memoryset[1].label == None
+
+
 def test_open_memoryset(readonly_memoryset, hf_dataset):
     fetched_memoryset = LabeledMemoryset.open(readonly_memoryset.name)
     assert fetched_memoryset is not None
@@ -96,9 +167,10 @@ def test_open_memoryset(readonly_memoryset, hf_dataset):
     assert fetched_memoryset.index_params == {"n_lists": 100}
 
 
-def test_open_memoryset_unauthenticated(unauthenticated, readonly_memoryset):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        LabeledMemoryset.open(readonly_memoryset.name)
+def test_open_memoryset_unauthenticated(unauthenticated_client, readonly_memoryset):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            LabeledMemoryset.open(readonly_memoryset.name)
 
 
 def test_open_memoryset_not_found():
@@ -111,9 +183,10 @@ def test_open_memoryset_invalid_input():
         LabeledMemoryset.open("not valid id")
 
 
-def test_open_memoryset_unauthorized(unauthorized, readonly_memoryset):
-    with pytest.raises(LookupError):
-        LabeledMemoryset.open(readonly_memoryset.name)
+def test_open_memoryset_unauthorized(unauthorized_client, readonly_memoryset):
+    with unauthorized_client.use():
+        with pytest.raises(LookupError):
+            LabeledMemoryset.open(readonly_memoryset.name)
 
 
 def test_all_memorysets(readonly_memoryset: LabeledMemoryset):
@@ -142,18 +215,21 @@ def test_all_memorysets_hidden(
     assert hidden_memoryset in all_memorysets
 
 
-def test_all_memorysets_unauthenticated(unauthenticated):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        LabeledMemoryset.all()
+def test_all_memorysets_unauthenticated(unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            LabeledMemoryset.all()
 
 
-def test_all_memorysets_unauthorized(unauthorized, readonly_memoryset):
-    assert readonly_memoryset not in LabeledMemoryset.all()
+def test_all_memorysets_unauthorized(unauthorized_client, readonly_memoryset):
+    with unauthorized_client.use():
+        assert readonly_memoryset not in LabeledMemoryset.all()
 
 
-def test_drop_memoryset_unauthenticated(unauthenticated, readonly_memoryset):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        LabeledMemoryset.drop(readonly_memoryset.name)
+def test_drop_memoryset_unauthenticated(unauthenticated_client, readonly_memoryset):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            LabeledMemoryset.drop(readonly_memoryset.name)
 
 
 def test_drop_memoryset_not_found():
@@ -163,9 +239,10 @@ def test_drop_memoryset_not_found():
     LabeledMemoryset.drop(str(uuid4()), if_not_exists="ignore")
 
 
-def test_drop_memoryset_unauthorized(unauthorized, readonly_memoryset):
-    with pytest.raises(LookupError):
-        LabeledMemoryset.drop(readonly_memoryset.name)
+def test_drop_memoryset_unauthorized(unauthorized_client, readonly_memoryset):
+    with unauthorized_client.use():
+        with pytest.raises(LookupError):
+            LabeledMemoryset.drop(readonly_memoryset.name)
 
 
 def test_update_memoryset_attributes(writable_memoryset: LabeledMemoryset):
@@ -304,6 +381,143 @@ def test_query_memoryset_with_feedback_metrics_sort(classification_model: Classi
     assert memories[-1].feedback_metrics["positive"]["avg"] == -1.0
 
 
+def test_labeled_memory_predictions_property(classification_model: ClassificationModel):
+    """Test that LabeledMemory.predictions() only returns classification predictions."""
+    # Given: A classification model with memories
+    memories = classification_model.memoryset.query(limit=1)
+    assert len(memories) > 0
+    memory = memories[0]
+
+    # When: I call the predictions method
+    predictions = memory.predictions()
+
+    # Then: It should return a list of ClassificationPrediction objects
+    assert isinstance(predictions, list)
+    for prediction in predictions:
+        assert prediction.__class__.__name__ == "ClassificationPrediction"
+        assert hasattr(prediction, "label")
+        assert not hasattr(prediction, "score") or prediction.score is None
+
+
+def test_scored_memory_predictions_property(regression_model: RegressionModel):
+    """Test that ScoredMemory.predictions() only returns regression predictions."""
+    # Given: A regression model with memories
+    memories = regression_model.memoryset.query(limit=1)
+    assert len(memories) > 0
+    memory = memories[0]
+
+    # When: I call the predictions method
+    predictions = memory.predictions()
+
+    # Then: It should return a list of RegressionPrediction objects
+    assert isinstance(predictions, list)
+    for prediction in predictions:
+        assert prediction.__class__.__name__ == "RegressionPrediction"
+        assert hasattr(prediction, "score")
+        assert not hasattr(prediction, "label") or prediction.label is None
+
+
+def test_memory_feedback_property(classification_model: ClassificationModel):
+    """Test that memory.feedback() returns feedback from relevant predictions."""
+    # Given: A prediction with recorded feedback
+    prediction = classification_model.predict("Test feedback")
+    feedback_category = f"test_feedback_{random.randint(0, 1000000)}"
+    prediction.record_feedback(category=feedback_category, value=True)
+
+    # And: A memory that was used in the prediction
+    memory_lookups = prediction.memory_lookups
+    assert len(memory_lookups) > 0
+    memory = memory_lookups[0]
+
+    # When: I access the feedback property
+    feedback = memory.feedback()
+
+    # Then: It should return feedback aggregated by category as a dict
+    assert isinstance(feedback, dict)
+    assert feedback_category in feedback
+    # Feedback values are lists (you may want to look at mean on the raw data)
+    assert isinstance(feedback[feedback_category], list)
+    assert len(feedback[feedback_category]) > 0
+    # For binary feedback, values should be booleans
+    assert isinstance(feedback[feedback_category][0], bool)
+
+
+def test_memory_predictions_method_parameters(classification_model: ClassificationModel):
+    """Test that memory.predictions() method supports pagination, sorting, and filtering."""
+    # Given: A classification model with memories
+    memories = classification_model.memoryset.query(limit=1)
+    assert len(memories) > 0
+    memory = memories[0]
+
+    # When: I call predictions with limit parameter
+    predictions_limited = memory.predictions(limit=2)
+
+    # Then: It should respect the limit
+    assert isinstance(predictions_limited, list)
+    assert len(predictions_limited) <= 2
+
+    # When: I call predictions with offset parameter
+    all_predictions = memory.predictions(limit=100)
+    if len(all_predictions) > 1:
+        predictions_offset = memory.predictions(limit=1, offset=1)
+        # Then: offset should skip the first prediction
+        assert predictions_offset[0].prediction_id != all_predictions[0].prediction_id
+
+    # When: I call predictions with sort parameter
+    predictions_sorted = memory.predictions(limit=10, sort=[("timestamp", "desc")])
+    # Then: It should return predictions (sorting verified by API)
+    assert isinstance(predictions_sorted, list)
+
+    # When: I call predictions with expected_label_match parameter
+    correct_predictions = memory.predictions(expected_label_match=True)
+    incorrect_predictions = memory.predictions(expected_label_match=False)
+    # Then: Both should return lists (correctness verified by API filtering)
+    assert isinstance(correct_predictions, list)
+    assert isinstance(incorrect_predictions, list)
+
+
+def test_memory_predictions_expected_label_filter(classification_model: ClassificationModel):
+    """Test that memory.predictions(expected_label_match=...) filters predictions by correctness."""
+    # Given: Make an initial prediction to learn the model's label for a known input
+    baseline_prediction = classification_model.predict("Filter test input", save_telemetry="sync")
+    original_label = baseline_prediction.label
+    alternate_label = 0 if original_label else 1
+
+    # When: Make a second prediction with an intentionally incorrect expected label
+    mismatched_prediction = classification_model.predict(
+        "Filter test input",
+        expected_labels=alternate_label,
+        save_telemetry="sync",
+    )
+    mismatched_memory = mismatched_prediction.memory_lookups[0]
+
+    # Then: The prediction should show up when filtering for incorrect predictions
+    incorrect_predictions = mismatched_memory.predictions(expected_label_match=False)
+    assert any(pred.prediction_id == mismatched_prediction.prediction_id for pred in incorrect_predictions)
+
+    # Produce a correct prediction (predicted label matches expected label)
+    correct_prediction = classification_model.predict(
+        "Filter test input",
+        expected_labels=original_label,
+        save_telemetry="sync",
+    )
+
+    # Ensure we are inspecting a memory used by both correct and incorrect predictions
+    correct_lookup_ids = {lookup.memory_id for lookup in correct_prediction.memory_lookups}
+    if mismatched_memory.memory_id not in correct_lookup_ids:
+        shared_lookup = next(
+            (lookup for lookup in mismatched_prediction.memory_lookups if lookup.memory_id in correct_lookup_ids),
+            None,
+        )
+        assert shared_lookup is not None, "No shared memory lookup between correct and incorrect predictions"
+        mismatched_memory = shared_lookup
+
+    # And: The correct prediction should appear when filtering for correct predictions
+    correct_predictions = mismatched_memory.predictions(expected_label_match=True)
+    assert any(pred.prediction_id == correct_prediction.prediction_id for pred in correct_predictions)
+    assert all(pred.prediction_id != mismatched_prediction.prediction_id for pred in correct_predictions)
+
+
 def test_insert_memories(writable_memoryset: LabeledMemoryset):
     writable_memoryset.refresh()
     prev_length = writable_memoryset.length
@@ -327,6 +541,7 @@ def test_insert_memories(writable_memoryset: LabeledMemoryset):
 
 
 @skip_in_prod("Production memorysets do not have session consistency guarantees")
+@skip_in_ci("CI environment may not have session consistency guarantees")
 def test_update_memories(writable_memoryset: LabeledMemoryset, hf_dataset: Dataset):
     # We've combined the update tests into one to avoid multiple expensive requests for a writable_memoryset
 
@@ -385,17 +600,6 @@ def test_clone_memoryset(readonly_memoryset: LabeledMemoryset):
     assert cloned_memoryset.insertion_status == Status.COMPLETED
 
 
-def test_embedding_evaluation(eval_datasource: Datasource):
-    results = LabeledMemoryset.run_embedding_evaluation(
-        eval_datasource, embedding_models=["CDE_SMALL"], neighbor_count=3
-    )
-    assert isinstance(results, list)
-    assert len(results) == 1
-    assert results[0] is not None
-    assert results[0]["embedding_model_name"] == "CDE_SMALL"
-    assert results[0]["embedding_model_path"] == "OrcaDB/cde-small-v1"
-
-
 @pytest.fixture(scope="function")
 async def test_group_potential_duplicates(writable_memoryset: LabeledMemoryset):
     writable_memoryset.insert(
@@ -492,7 +696,8 @@ def test_scored_memoryset(scored_memoryset: ScoredMemoryset):
     assert isinstance(scored_memoryset[0], ScoredMemory)
     assert scored_memoryset[0].value == "i love soup"
     assert scored_memoryset[0].score is not None
-    assert scored_memoryset[0].metadata == {"key": "g1", "source_id": "s1", "label": 0}
+    assert scored_memoryset[0].metadata == {"key": "g1", "label": 0}
+    assert scored_memoryset[0].source_id == "s1"
     lookup = scored_memoryset.search("i love soup", count=1)
     assert len(lookup) == 1
     assert lookup[0].score is not None
@@ -508,3 +713,80 @@ def test_update_scored_memory(scored_memoryset: ScoredMemoryset):
     assert scored_memoryset[0].label == 3
     memory.update(label=4)
     assert scored_memoryset[0].label == 4
+
+
+@pytest.mark.asyncio
+async def test_insert_memories_async_single(writable_memoryset: LabeledMemoryset):
+    """Test async insertion of a single memory"""
+    await writable_memoryset.arefresh()
+    prev_length = writable_memoryset.length
+
+    await writable_memoryset.ainsert(dict(value="async tomato soup is my favorite", label=0, key="async_test"))
+
+    await writable_memoryset.arefresh()
+    assert writable_memoryset.length == prev_length + 1
+    last_memory = writable_memoryset[-1]
+    assert last_memory.value == "async tomato soup is my favorite"
+    assert last_memory.label == 0
+    assert last_memory.metadata["key"] == "async_test"
+
+
+@pytest.mark.asyncio
+async def test_insert_memories_async_batch(writable_memoryset: LabeledMemoryset):
+    """Test async insertion of multiple memories"""
+    await writable_memoryset.arefresh()
+    prev_length = writable_memoryset.length
+
+    await writable_memoryset.ainsert(
+        [
+            dict(value="async batch soup is delicious", label=0, key="batch_test_1"),
+            dict(value="async batch cats are adorable", label=1, key="batch_test_2"),
+        ]
+    )
+
+    await writable_memoryset.arefresh()
+    assert writable_memoryset.length == prev_length + 2
+
+    # Check the inserted memories
+    last_two_memories = writable_memoryset[-2:]
+    values = [memory.value for memory in last_two_memories]
+    labels = [memory.label for memory in last_two_memories]
+    keys = [memory.metadata.get("key") for memory in last_two_memories]
+
+    assert "async batch soup is delicious" in values
+    assert "async batch cats are adorable" in values
+    assert 0 in labels
+    assert 1 in labels
+    assert "batch_test_1" in keys
+    assert "batch_test_2" in keys
+
+
+@pytest.mark.asyncio
+async def test_insert_memories_async_with_source_id(writable_memoryset: LabeledMemoryset):
+    """Test async insertion with source_id and metadata"""
+    await writable_memoryset.arefresh()
+    prev_length = writable_memoryset.length
+
+    await writable_memoryset.ainsert(
+        dict(
+            value="async soup with source id", label=0, source_id="async_source_123", custom_field="async_custom_value"
+        )
+    )
+
+    await writable_memoryset.arefresh()
+    assert writable_memoryset.length == prev_length + 1
+    last_memory = writable_memoryset[-1]
+    assert last_memory.value == "async soup with source id"
+    assert last_memory.label == 0
+    assert last_memory.source_id == "async_source_123"
+    assert last_memory.metadata["custom_field"] == "async_custom_value"
+
+
+@pytest.mark.asyncio
+async def test_insert_memories_async_unauthenticated(
+    unauthenticated_async_client, writable_memoryset: LabeledMemoryset
+):
+    """Test async insertion with invalid authentication"""
+    with unauthenticated_async_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            await writable_memoryset.ainsert(dict(value="this should fail", label=0))