orca-sdk 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

orca_sdk/memoryset_test.py

@@ -8,7 +8,14 @@ from .classification_model import ClassificationModel
 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 .memoryset import (
+    LabeledMemory,
+    LabeledMemoryset,
+    ScoredMemory,
+    ScoredMemoryset,
+    Status,
+)
+from .regression_model import RegressionModel
 
 """
 Test Performance Note:
@@ -112,6 +119,51 @@ def test_if_exists_open_reuses_existing_datasource(
     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 is None
+    assert memoryset[1].label is None
+
+
 def test_open_memoryset(readonly_memoryset, hf_dataset):
     fetched_memoryset = LabeledMemoryset.open(readonly_memoryset.name)
     assert fetched_memoryset is not None
@@ -239,6 +291,100 @@ def test_search_count(readonly_memoryset: LabeledMemoryset):
     assert memory_lookups[2].label == 0
 
 
+def test_search_with_partition_id(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Search within a specific partition - use "soup" which appears in both p1 and p2
+    # Use exclude_global to ensure we only get results from the specified partition
+    memory_lookups = readonly_partitioned_memoryset.search(
+        "soup", partition_id="p1", partition_filter_mode="exclude_global", count=5
+    )
+    assert len(memory_lookups) > 0
+    # All results should be from partition p1 when partition_id is specified
+    assert all(
+        memory.partition_id == "p1" for memory in memory_lookups
+    ), f"Expected all results from partition p1, but got: {[m.partition_id for m in memory_lookups]}"
+
+    # Search in a different partition - use "cats" which appears in both p1 and p2
+    memory_lookups_p2 = readonly_partitioned_memoryset.search(
+        "cats", partition_id="p2", partition_filter_mode="exclude_global", count=5
+    )
+    assert len(memory_lookups_p2) > 0
+    # All results should be from partition p2 when partition_id is specified
+    assert all(
+        memory.partition_id == "p2" for memory in memory_lookups_p2
+    ), f"Expected all results from partition p2, but got: {[m.partition_id for m in memory_lookups_p2]}"
+
+
+def test_search_with_partition_filter_mode_exclude_global(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Search excluding global memories - need to specify a partition_id when using exclude_global
+    # This tests that exclude_global works with a specific partition
+    memory_lookups = readonly_partitioned_memoryset.search(
+        "soup", partition_id="p1", partition_filter_mode="exclude_global", count=5
+    )
+    assert len(memory_lookups) > 0
+    # All results should have a partition_id (not None) and be from p1
+    assert all(memory.partition_id == "p1" for memory in memory_lookups)
+
+
+def test_search_with_partition_filter_mode_only_global(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Search only in global memories (partition_id=None in the data)
+    # Use a query that matches global memories and a reasonable count
+    memory_lookups = readonly_partitioned_memoryset.search("beach", partition_filter_mode="only_global", count=3)
+    # Should get at least some results (may be fewer than requested if not enough global memories match)
+    assert len(memory_lookups) > 0
+    # All results should be global (partition_id is None)
+    partition_ids = {memory.partition_id for memory in memory_lookups}
+    # When using only_global, all results should be global (either None)
+    assert all(
+        memory.partition_id is None for memory in memory_lookups
+    ), f"Expected all results to be global (partition_id=None), but got partition_ids: {partition_ids}"
+
+
+def test_search_with_partition_filter_mode_include_global(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Search including global memories (default behavior)
+    # Use a reasonable count that won't exceed available memories
+    memory_lookups = readonly_partitioned_memoryset.search(
+        "i love soup", partition_filter_mode="include_global", count=5
+    )
+    assert len(memory_lookups) > 0
+    # Results can include both partitioned and global memories
+    partition_ids = {memory.partition_id for memory in memory_lookups}
+    # Should have at least one partition or global memory
+    assert len(partition_ids) > 0
+
+
+def test_search_with_partition_filter_mode_ignore_partitions(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Search ignoring partition filtering entirely
+    memory_lookups = readonly_partitioned_memoryset.search(
+        "i love soup", partition_filter_mode="ignore_partitions", count=10
+    )
+    assert len(memory_lookups) > 0
+    # Results can come from any partition or global
+    partition_ids = {memory.partition_id for memory in memory_lookups}
+    # Should have results from multiple partitions/global
+    assert len(partition_ids) >= 1
+
+
+def test_search_multiple_queries_with_partition_id(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Search multiple queries within a specific partition
+    memory_lookups = readonly_partitioned_memoryset.search(["i love soup", "cats are cute"], partition_id="p1", count=3)
+    assert len(memory_lookups) == 2
+    assert len(memory_lookups[0]) > 0
+    assert len(memory_lookups[1]) > 0
+    # All results should be from partition p1
+    assert all(memory.partition_id == "p1" for memory in memory_lookups[0])
+    assert all(memory.partition_id == "p1" for memory in memory_lookups[1])
+
+
+def test_search_with_partition_id_and_filter_mode(readonly_partitioned_memoryset: LabeledMemoryset):
+    # When partition_id is specified, partition_filter_mode should still work
+    # Search in p1 with exclude_global (should only return p1 results)
+    memory_lookups = readonly_partitioned_memoryset.search(
+        "i love soup", partition_id="p1", partition_filter_mode="exclude_global", count=5
+    )
+    assert len(memory_lookups) > 0
+    assert all(memory.partition_id == "p1" for memory in memory_lookups)
+
+
 def test_get_memory_at_index(readonly_memoryset: LabeledMemoryset, hf_dataset: Dataset, label_names: list[str]):
     memory = readonly_memoryset[0]
     assert memory.value == hf_dataset[0]["value"]
@@ -335,6 +481,292 @@ def test_query_memoryset_with_feedback_metrics_sort(classification_model: Classi
     assert memories[-1].feedback_metrics["positive"]["avg"] == -1.0
 
 
+def test_query_memoryset_with_partition_id(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query with partition_id and include_global (default) - includes both p1 and global memories
+    memories = readonly_partitioned_memoryset.query(partition_id="p1")
+    assert len(memories) == 15  # 8 p1 + 7 global = 15
+    # Results should include both p1 and global memories
+    partition_ids = {memory.partition_id for memory in memories}
+    assert "p1" in partition_ids
+    assert None in partition_ids
+
+
+def test_query_memoryset_with_partition_id_and_exclude_global(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query with partition_id and exclude_global mode - only returns p1 memories
+    memories = readonly_partitioned_memoryset.query(partition_id="p1", partition_filter_mode="exclude_global")
+    assert len(memories) == 8  # Only 8 p1 memories (no global)
+    # All results should be from partition p1 (no global memories)
+    assert all(memory.partition_id == "p1" for memory in memories)
+
+
+def test_query_memoryset_with_partition_id_and_include_global(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query with partition_id and include_global mode (default) - includes both p1 and global
+    memories = readonly_partitioned_memoryset.query(partition_id="p1", partition_filter_mode="include_global")
+    assert len(memories) == 15  # 8 p1 + 7 global = 15
+    # Results should include both p1 and global memories
+    partition_ids = {memory.partition_id for memory in memories}
+    assert "p1" in partition_ids
+    assert None in partition_ids
+
+
+def test_query_memoryset_with_partition_filter_mode_exclude_global(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query excluding global memories requires a partition_id
+    # Test with a specific partition_id
+    memories = readonly_partitioned_memoryset.query(partition_id="p1", partition_filter_mode="exclude_global")
+    assert len(memories) == 8  # Only p1 memories
+    # All results should have a partition_id (not global)
+    assert all(memory.partition_id == "p1" for memory in memories)
+
+
+def test_query_memoryset_with_partition_filter_mode_only_global(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query only in global memories
+    memories = readonly_partitioned_memoryset.query(partition_filter_mode="only_global")
+    assert len(memories) == 7  # There are 7 global memories in SAMPLE_DATA
+    # All results should be global (partition_id is None)
+    assert all(memory.partition_id is None for memory in memories)
+
+
+def test_query_memoryset_with_partition_filter_mode_include_global(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query including global memories - when no partition_id is specified,
+    # include_global seems to only return global memories
+    memories = readonly_partitioned_memoryset.query(partition_filter_mode="include_global")
+    # Based on actual behavior, this returns only global memories
+    assert len(memories) == 7
+    # All results should be global
+    assert all(memory.partition_id is None for memory in memories)
+
+
+def test_query_memoryset_with_partition_filter_mode_ignore_partitions(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query ignoring partition filtering entirely - returns all memories
+    memories = readonly_partitioned_memoryset.query(partition_filter_mode="ignore_partitions", limit=100)
+    assert len(memories) == 22  # All 22 memories
+    # Results can come from any partition or global
+    partition_ids = {memory.partition_id for memory in memories}
+    # Should have results from multiple partitions/global
+    assert len(partition_ids) >= 1
+    # Verify we have p1, p2, and global
+    assert "p1" in partition_ids
+    assert "p2" in partition_ids
+    assert None in partition_ids
+
+
+def test_query_memoryset_with_filters_and_partition_id(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query with filters and partition_id
+    memories = readonly_partitioned_memoryset.query(filters=[("label", "==", 0)], partition_id="p1")
+    assert len(memories) > 0
+    # All results should match the filter and be from partition p1
+    assert all(memory.label == 0 for memory in memories)
+    assert all(memory.partition_id == "p1" for memory in memories)
+
+
+def test_query_memoryset_with_filters_and_partition_filter_mode(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query with filters and partition_filter_mode - exclude_global requires partition_id
+    memories = readonly_partitioned_memoryset.query(
+        filters=[("label", "==", 1)], partition_id="p1", partition_filter_mode="exclude_global"
+    )
+    assert len(memories) > 0
+    # All results should match the filter and be from p1 (not global)
+    assert all(memory.label == 1 for memory in memories)
+    assert all(memory.partition_id == "p1" for memory in memories)
+
+
+def test_query_memoryset_with_limit_and_partition_id(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query with limit and partition_id
+    memories = readonly_partitioned_memoryset.query(partition_id="p2", limit=3)
+    assert len(memories) == 3
+    # All results should be from partition p2
+    assert all(memory.partition_id == "p2" for memory in memories)
+
+
+def test_query_memoryset_with_offset_and_partition_id(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query with offset and partition_id - use exclude_global to get only p1 memories
+    memories_page1 = readonly_partitioned_memoryset.query(
+        partition_id="p1", partition_filter_mode="exclude_global", limit=5
+    )
+    memories_page2 = readonly_partitioned_memoryset.query(
+        partition_id="p1", partition_filter_mode="exclude_global", offset=5, limit=5
+    )
+    assert len(memories_page1) == 5
+    assert len(memories_page2) == 3  # Only 3 remaining p1 memories (8 total - 5 = 3)
+    # All results should be from partition p1
+    assert all(memory.partition_id == "p1" for memory in memories_page1)
+    assert all(memory.partition_id == "p1" for memory in memories_page2)
+    # Results should be different (pagination works)
+    memory_ids_page1 = {memory.memory_id for memory in memories_page1}
+    memory_ids_page2 = {memory.memory_id for memory in memories_page2}
+    assert memory_ids_page1.isdisjoint(memory_ids_page2)
+
+
+def test_query_memoryset_with_partition_id_p2(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query a different partition to verify it works
+    # With include_global (default), it includes both p2 and global memories
+    memories = readonly_partitioned_memoryset.query(partition_id="p2")
+    assert len(memories) == 14  # 7 p2 + 7 global = 14
+    # Results should include both p2 and global memories
+    partition_ids = {memory.partition_id for memory in memories}
+    assert "p2" in partition_ids
+    assert None in partition_ids
+
+
+def test_query_memoryset_with_metadata_filter_and_partition_id(readonly_partitioned_memoryset: LabeledMemoryset):
+    # Query with metadata filter and partition_id
+    memories = readonly_partitioned_memoryset.query(filters=[("metadata.key", "==", "g1")], partition_id="p1")
+    assert len(memories) > 0
+    # All results should match the metadata filter and be from partition p1
+    assert all(memory.metadata.get("key") == "g1" for memory in memories)
+    assert all(memory.partition_id == "p1" for memory in memories)
+
+
+def test_query_memoryset_with_partition_filter_mode_only_global_and_filters(
+    readonly_partitioned_memoryset: LabeledMemoryset,
+):
+    # Query only global memories with filters
+    memories = readonly_partitioned_memoryset.query(
+        filters=[("metadata.key", "==", "g3")], partition_filter_mode="only_global"
+    )
+    assert len(memories) > 0
+    # All results should match the filter and be global
+    assert all(memory.metadata.get("key") == "g3" for memory in memories)
+    assert all(memory.partition_id is None for memory in memories)
+
+
+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
@@ -513,7 +945,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, "partition_id": "p1"}
+    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