orca-sdk 0.0.94__py3-none-any.whl → 0.0.96__py3-none-any.whl

orca_sdk/memoryset_test.py

@@ -1,19 +1,15 @@
+import os
 import random
-import time
-from typing import Generator
 from uuid import uuid4
 
 import pytest
-from datasets import ClassLabel, Features, Value
 from datasets.arrow_dataset import Dataset
 
-from orca_sdk.conftest import SAMPLE_DATA
-
-from ._generated_api_client.models import CascadingEditSuggestion
 from .classification_model import ClassificationModel
+from .conftest import skip_in_prod
 from .datasource import Datasource
 from .embedding_model import PretrainedEmbeddingModel
-from .memoryset import LabeledMemoryset, TaskStatus
+from .memoryset import LabeledMemoryset, ScoredMemory, ScoredMemoryset, Status
 
 """
 Test Performance Note:
@@ -37,9 +33,11 @@ def test_create_memoryset(readonly_memoryset: LabeledMemoryset, hf_dataset: Data
     assert readonly_memoryset.name == "test_readonly_memoryset"
     assert readonly_memoryset.embedding_model == PretrainedEmbeddingModel.GTE_BASE
     assert readonly_memoryset.label_names == label_names
-    assert readonly_memoryset.insertion_status == TaskStatus.COMPLETED
+    assert readonly_memoryset.insertion_status == Status.COMPLETED
     assert isinstance(readonly_memoryset.length, int)
     assert readonly_memoryset.length == len(hf_dataset)
+    assert readonly_memoryset.index_type == "IVF_FLAT"
+    assert readonly_memoryset.index_params == {"n_lists": 100}
 
 
 def test_create_memoryset_unauthenticated(unauthenticated, datasource):
@@ -95,6 +93,8 @@ def test_open_memoryset(readonly_memoryset, hf_dataset):
     assert fetched_memoryset is not None
     assert fetched_memoryset.name == readonly_memoryset.name
     assert fetched_memoryset.length == len(hf_dataset)
+    assert fetched_memoryset.index_type == "IVF_FLAT"
+    assert fetched_memoryset.index_params == {"n_lists": 100}
 
 
 def test_open_memoryset_unauthenticated(unauthenticated, readonly_memoryset):
@@ -149,15 +149,25 @@ def test_drop_memoryset_unauthorized(unauthorized, readonly_memoryset):
         LabeledMemoryset.drop(readonly_memoryset.name)
 
 
-def test_update_memoryset_metadata(writable_memoryset: LabeledMemoryset):
-    # NOTE: We're combining multiple tests into one here to avoid multiple API calls
-
-    writable_memoryset.update_metadata(description="New description")
+def test_update_memoryset_attributes(writable_memoryset: LabeledMemoryset):
+    original_label_names = writable_memoryset.label_names
+    writable_memoryset.set(description="New description")
     assert writable_memoryset.description == "New description"
 
-    writable_memoryset.update_metadata(description=None)
+    writable_memoryset.set(description=None)
     assert writable_memoryset.description is None
 
+    writable_memoryset.set(name="New_name")
+    assert writable_memoryset.name == "New_name"
+
+    writable_memoryset.set(name="test_writable_memoryset")
+    assert writable_memoryset.name == "test_writable_memoryset"
+
+    assert writable_memoryset.label_names == original_label_names
+
+    writable_memoryset.set(label_names=["New label 1", "New label 2"])
+    assert writable_memoryset.label_names == ["New label 1", "New label 2"]
+
 
 def test_search(readonly_memoryset: LabeledMemoryset):
     memory_lookups = readonly_memoryset.search(["i love soup", "cats are cute"])
@@ -214,11 +224,11 @@ def test_query_memoryset(readonly_memoryset: LabeledMemoryset):
     assert len(memories) == 8
     assert all(memory.label == 1 for memory in memories)
     assert len(readonly_memoryset.query(limit=2)) == 2
-    assert len(readonly_memoryset.query(filters=[("metadata.key", "==", "val1")])) == 1
+    assert len(readonly_memoryset.query(filters=[("metadata.key", "==", "g2")])) == 4
 
 
-def test_query_memoryset_with_feedback_metrics(model: ClassificationModel):
-    prediction = model.predict("Do you love soup?")
+def test_query_memoryset_with_feedback_metrics(classification_model: ClassificationModel):
+    prediction = classification_model.predict("Do you love soup?")
     feedback_name = f"correct_{random.randint(0, 1000000)}"
     prediction.record_feedback(category=feedback_name, value=prediction.label == 0)
     memories = prediction.memoryset.query(filters=[("label", "==", 0)], with_feedback_metrics=True)
@@ -239,8 +249,8 @@ def test_query_memoryset_with_feedback_metrics(model: ClassificationModel):
         assert isinstance(memory.lookup_count, int)
 
 
-def test_query_memoryset_with_feedback_metrics_filter(model: ClassificationModel):
-    prediction = model.predict("Do you love soup?")
+def test_query_memoryset_with_feedback_metrics_filter(classification_model: ClassificationModel):
+    prediction = classification_model.predict("Do you love soup?")
     prediction.record_feedback(category="accurate", value=prediction.label == 0)
     memories = prediction.memoryset.query(
         filters=[("feedback_metrics.accurate.avg", ">", 0.5)], with_feedback_metrics=True
@@ -254,10 +264,10 @@ def test_query_memoryset_with_feedback_metrics_filter(model: ClassificationModel
         assert memory.feedback_metrics["accurate"]["count"] == 1
 
 
-def test_query_memoryset_with_feedback_metrics_sort(model: ClassificationModel):
-    prediction = model.predict("Do you love soup?")
+def test_query_memoryset_with_feedback_metrics_sort(classification_model: ClassificationModel):
+    prediction = classification_model.predict("Do you love soup?")
     prediction.record_feedback(category="positive", value=1.0)
-    prediction2 = model.predict("Do you like cats?")
+    prediction2 = classification_model.predict("Do you like cats?")
     prediction2.record_feedback(category="positive", value=-1.0)
 
     memories = prediction.memoryset.query(
@@ -294,6 +304,7 @@ def test_insert_memories(writable_memoryset: LabeledMemoryset):
     assert last_memory.source_id == "test"
 
 
+@skip_in_prod("Production memorysets do 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
 
@@ -302,6 +313,7 @@ def test_update_memories(writable_memoryset: LabeledMemoryset, hf_dataset: Datas
     updated_memory = writable_memoryset.update(dict(memory_id=memory_id, value="i love soup so much"))
     assert updated_memory.value == "i love soup so much"
     assert updated_memory.label == hf_dataset[0]["label"]
+    writable_memoryset.refresh()  # Refresh to ensure consistency after update
     assert writable_memoryset.get(memory_id).value == "i love soup so much"
 
     # test updating a memory instance
@@ -348,7 +360,7 @@ def test_clone_memoryset(readonly_memoryset: LabeledMemoryset):
     assert cloned_memoryset.name == "test_cloned_memoryset"
     assert cloned_memoryset.length == readonly_memoryset.length
     assert cloned_memoryset.embedding_model == PretrainedEmbeddingModel.DISTILBERT
-    assert cloned_memoryset.insertion_status == TaskStatus.COMPLETED
+    assert cloned_memoryset.insertion_status == Status.COMPLETED
 
 
 def test_embedding_evaluation(eval_datasource: Datasource):
@@ -363,7 +375,6 @@ def test_embedding_evaluation(eval_datasource: Datasource):
     assert response["evaluation_results"][0] is not None
     assert response["evaluation_results"][0]["embedding_model_name"] == "CDE_SMALL"
     assert response["evaluation_results"][0]["embedding_model_path"] == "OrcaDB/cde-small-v1"
-    Datasource.drop("eval_datasource")
 
 
 @pytest.fixture(scope="function")
@@ -455,3 +466,25 @@ def test_drop_memoryset(writable_memoryset: LabeledMemoryset):
     assert LabeledMemoryset.exists(writable_memoryset.name)
     LabeledMemoryset.drop(writable_memoryset.name)
     assert not LabeledMemoryset.exists(writable_memoryset.name)
+
+
+def test_scored_memoryset(scored_memoryset: ScoredMemoryset):
+    assert scored_memoryset.length == 16
+    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}
+    lookup = scored_memoryset.search("i love soup", count=1)
+    assert len(lookup) == 1
+    assert lookup[0].score < 0.11
+
+
+@skip_in_prod("Production memorysets do not have session consistency guarantees")
+def test_update_scored_memory(scored_memoryset: ScoredMemoryset):
+    # we are only updating an inconsequential metadata field so that we don't affect other tests
+    memory = scored_memoryset[0]
+    assert memory.label == 0
+    scored_memoryset.update(dict(memory_id=memory.memory_id, label=3))
+    assert scored_memoryset[0].label == 3
+    memory.update(label=4)
+    assert scored_memoryset[0].label == 4