orca-sdk 0.0.92__py3-none-any.whl → 0.0.93__py3-none-any.whl

orca_sdk/conftest.py

@@ -17,6 +17,8 @@ logging.basicConfig(level=logging.INFO)
 
 os.environ["ORCA_API_URL"] = os.environ.get("ORCA_API_URL", "http://localhost:1584/")
 
+os.environ["ORCA_SAVE_TELEMETRY_SYNCHRONOUSLY"] = "true"
+
 
 def _create_org_id():
     # UUID start to identify test data (0xtest...)
@@ -69,22 +71,22 @@ def label_names():
 
 
 SAMPLE_DATA = [
-    {"text": "i love soup", "label": 0, "key": "val1", "score": 0.1, "source_id": "s1"},
-    {"text": "cats are cute", "label": 1, "key": "val2", "score": 0.2, "source_id": "s2"},
-    {"text": "soup is good", "label": 0, "key": "val3", "score": 0.3, "source_id": "s3"},
-    {"text": "i love cats", "label": 1, "key": "val4", "score": 0.4, "source_id": "s4"},
-    {"text": "everyone loves cats", "label": 1, "key": "val5", "score": 0.5, "source_id": "s5"},
-    {"text": "soup is great for the winter", "label": 0, "key": "val6", "score": 0.6, "source_id": "s6"},
-    {"text": "hot soup on a rainy day!", "label": 0, "key": "val7", "score": 0.7, "source_id": "s7"},
-    {"text": "cats sleep all day", "label": 1, "key": "val8", "score": 0.8, "source_id": "s8"},
-    {"text": "homemade soup recipes", "label": 0, "key": "val9", "score": 0.9, "source_id": "s9"},
-    {"text": "cats purr when happy", "label": 1, "key": "val10", "score": 1.0, "source_id": "s10"},
-    {"text": "chicken noodle soup is classic", "label": 0, "key": "val11", "score": 1.1, "source_id": "s11"},
-    {"text": "kittens are baby cats", "label": 1, "key": "val12", "score": 1.2, "source_id": "s12"},
-    {"text": "soup can be served cold too", "label": 0, "key": "val13", "score": 1.3, "source_id": "s13"},
-    {"text": "cats have nine lives", "label": 1, "key": "val14", "score": 1.4, "source_id": "s14"},
-    {"text": "tomato soup with grilled cheese", "label": 0, "key": "val15", "score": 1.5, "source_id": "s15"},
-    {"text": "cats are independent animals", "label": 1, "key": "val16", "score": 1.6, "source_id": "s16"},
+    {"value": "i love soup", "label": 0, "key": "val1", "score": 0.1, "source_id": "s1"},
+    {"value": "cats are cute", "label": 1, "key": "val2", "score": 0.2, "source_id": "s2"},
+    {"value": "soup is good", "label": 0, "key": "val3", "score": 0.3, "source_id": "s3"},
+    {"value": "i love cats", "label": 1, "key": "val4", "score": 0.4, "source_id": "s4"},
+    {"value": "everyone loves cats", "label": 1, "key": "val5", "score": 0.5, "source_id": "s5"},
+    {"value": "soup is great for the winter", "label": 0, "key": "val6", "score": 0.6, "source_id": "s6"},
+    {"value": "hot soup on a rainy day!", "label": 0, "key": "val7", "score": 0.7, "source_id": "s7"},
+    {"value": "cats sleep all day", "label": 1, "key": "val8", "score": 0.8, "source_id": "s8"},
+    {"value": "homemade soup recipes", "label": 0, "key": "val9", "score": 0.9, "source_id": "s9"},
+    {"value": "cats purr when happy", "label": 1, "key": "val10", "score": 1.0, "source_id": "s10"},
+    {"value": "chicken noodle soup is classic", "label": 0, "key": "val11", "score": 1.1, "source_id": "s11"},
+    {"value": "kittens are baby cats", "label": 1, "key": "val12", "score": 1.2, "source_id": "s12"},
+    {"value": "soup can be served cold too", "label": 0, "key": "val13", "score": 1.3, "source_id": "s13"},
+    {"value": "cats have nine lives", "label": 1, "key": "val14", "score": 1.4, "source_id": "s14"},
+    {"value": "tomato soup with grilled cheese", "label": 0, "key": "val15", "score": 1.5, "source_id": "s15"},
+    {"value": "cats are independent animals", "label": 1, "key": "val16", "score": 1.6, "source_id": "s16"},
 ]
 
 
@@ -94,7 +96,7 @@ def hf_dataset(label_names):
         SAMPLE_DATA,
         features=Features(
             {
-                "text": Value("string"),
+                "value": Value("string"),
                 "label": ClassLabel(names=label_names),
                 "key": Value("string"),
                 "score": Value("float"),
@@ -106,23 +108,82 @@ def hf_dataset(label_names):
 
 @pytest.fixture(scope="session")
 def datasource(hf_dataset) -> Datasource:
-    return Datasource.from_hf_dataset("test_datasource", hf_dataset)
+    datasource = Datasource.from_hf_dataset("test_datasource", hf_dataset)
+    return datasource
+
+
+EVAL_DATASET = [
+    {"value": "chicken noodle soup is the best", "label": 1},
+    {"value": "cats are cute", "label": 0},
+    {"value": "soup is great for the winter", "label": 0},
+    {"value": "i love cats", "label": 1},
+]
 
 
 @pytest.fixture(scope="session")
-def memoryset(datasource) -> LabeledMemoryset:
-    return LabeledMemoryset.create(
-        "test_memoryset",
+def eval_datasource() -> Datasource:
+    eval_datasource = Datasource.from_list("eval_datasource", EVAL_DATASET)
+    return eval_datasource
+
+
+@pytest.fixture(scope="session")
+def eval_dataset() -> Dataset:
+    eval_dataset = Dataset.from_list(EVAL_DATASET)
+    return eval_dataset
+
+
+@pytest.fixture(scope="session")
+def readonly_memoryset(datasource: Datasource) -> LabeledMemoryset:
+    memoryset = LabeledMemoryset.create(
+        "test_readonly_memoryset",
         datasource=datasource,
         embedding_model=PretrainedEmbeddingModel.GTE_BASE,
-        value_column="text",
         source_id_column="source_id",
         max_seq_length_override=32,
     )
+    return memoryset
+
+
+@pytest.fixture(scope="function")
+def writable_memoryset(datasource: Datasource, api_key: str) -> Generator[LabeledMemoryset, None, None]:
+    """
+    Function-scoped fixture that provides a writable memoryset for tests that mutate state.
+
+    This fixture creates a fresh `LabeledMemoryset` named 'test_writable_memoryset' before each test.
+    After the test, it attempts to restore the memoryset to its initial state by deleting any added entries
+    and reinserting sample data — unless the memoryset has been dropped by the test itself, in which case
+    it will be recreated on the next invocation.
+
+    Note: Re-creating the memoryset from scratch is surprisingly more expensive than cleaning it up.
+    """
+    # It shouldn't be possible for this memoryset to already exist
+    memoryset = LabeledMemoryset.create(
+        "test_writable_memoryset",
+        datasource=datasource,
+        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
+        source_id_column="source_id",
+        max_seq_length_override=32,
+        if_exists="open",
+    )
+    try:
+        yield memoryset
+    finally:
+        # Restore the memoryset to a clean state for the next test.
+        OrcaCredentials.set_api_key(api_key, check_validity=False)
+
+        if LabeledMemoryset.exists("test_writable_memoryset"):
+            memory_ids = [memoryset[i].memory_id for i in range(len(memoryset))]
+
+            if memory_ids:
+                memoryset.delete(memory_ids)
+            assert len(memoryset) == 0
+            memoryset.insert(SAMPLE_DATA)
+        # If the test dropped the memoryset, do nothing — it will be recreated on the next use.
 
 
 @pytest.fixture(scope="session")
-def model(memoryset) -> ClassificationModel:
-    return ClassificationModel.create(
-        "test_model", memoryset, num_classes=2, memory_lookup_count=3, description="test_description"
+def model(readonly_memoryset: LabeledMemoryset) -> ClassificationModel:
+    model = ClassificationModel.create(
+        "test_model", readonly_memoryset, num_classes=2, memory_lookup_count=3, description="test_description"
     )
+    return model

orca_sdk/datasource.py

@@ -12,6 +12,7 @@ import pyarrow as pa
 from datasets import Dataset
 from torch.utils.data import DataLoader as TorchDataLoader
 from torch.utils.data import Dataset as TorchDataset
+from tqdm.auto import tqdm
 
 from ._generated_api_client.api import (
     delete_datasource,
@@ -25,6 +26,7 @@ from ._generated_api_client.client import get_client
 from ._generated_api_client.models import ColumnType, DatasourceMetadata
 from ._utils.common import CreateMode, DropMode
 from ._utils.data_parsing import hf_dataset_from_disk, hf_dataset_from_torch
+from ._utils.tqdm_file_reader import TqdmFileReader
 
 
 class Datasource:
@@ -113,19 +115,27 @@ class Datasource:
         with tempfile.TemporaryDirectory() as tmp_dir:
             dataset.save_to_disk(tmp_dir)
             files = []
-            for file_path in Path(tmp_dir).iterdir():
-                buffered_reader = open(file_path, "rb")
-                files.append(("files", buffered_reader))
-
-            # Do not use Generated client for this endpoint b/c it does not handle files properly
-            metadata = parse_create_response(
-                response=client.get_httpx_client().request(
-                    method="post",
-                    url="/datasource/",
-                    files=files,
-                    data={"name": name, "description": description},
+
+            # Calculate total size for all files
+            file_paths = list(Path(tmp_dir).iterdir())
+            total_size = sum(file_path.stat().st_size for file_path in file_paths)
+
+            with tqdm(total=total_size, unit="B", unit_scale=True, desc="Uploading") as pbar:
+                for file_path in file_paths:
+                    buffered_reader = open(file_path, "rb")
+                    tqdm_reader = TqdmFileReader(buffered_reader, pbar)
+                    files.append(("files", (file_path.name, tqdm_reader)))
+
+                # Do not use Generated client for this endpoint b/c it does not handle files properly
+                metadata = parse_create_response(
+                    response=client.get_httpx_client().request(
+                        method="post",
+                        url="/datasource/",
+                        files=files,
+                        data={"name": name, "description": description},
+                    )
                 )
-            )
+
             return cls(metadata=metadata)
 
     @classmethod

orca_sdk/embedding_model_test.py

@@ -53,7 +53,7 @@ def test_embed_text_unauthenticated(unauthenticated):
 
 @pytest.fixture(scope="session")
 def finetuned_model(datasource) -> FinetunedEmbeddingModel:
-    return PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model", datasource, value_column="text")
+    return PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model", datasource)
 
 
 def test_finetune_model_with_datasource(finetuned_model: FinetunedEmbeddingModel):
@@ -65,8 +65,10 @@ def test_finetune_model_with_datasource(finetuned_model: FinetunedEmbeddingModel
     assert finetuned_model._status == TaskStatus.COMPLETED
 
 
-def test_finetune_model_with_memoryset(memoryset: LabeledMemoryset):
-    finetuned_model = PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model_from_memoryset", memoryset)
+def test_finetune_model_with_memoryset(readonly_memoryset: LabeledMemoryset):
+    finetuned_model = PretrainedEmbeddingModel.DISTILBERT.finetune(
+        "test_finetuned_model_from_memoryset", readonly_memoryset
+    )
     assert finetuned_model is not None
     assert finetuned_model.name == "test_finetuned_model_from_memoryset"
     assert finetuned_model.base_model == PretrainedEmbeddingModel.DISTILBERT
@@ -109,7 +111,6 @@ def test_use_finetuned_model_in_memoryset(datasource: Datasource, finetuned_mode
         "test_memoryset_finetuned_model",
         datasource,
         embedding_model=finetuned_model,
-        value_column="text",
     )
     assert memoryset is not None
     assert memoryset.name == "test_memoryset_finetuned_model"
@@ -152,7 +153,7 @@ def test_all_finetuned_models_unauthorized(unauthorized, finetuned_model: Finetu
 
 
 def test_drop_finetuned_model(datasource: Datasource):
-    PretrainedEmbeddingModel.DISTILBERT.finetune("finetuned_model_to_delete", datasource, value_column="text")
+    PretrainedEmbeddingModel.DISTILBERT.finetune("finetuned_model_to_delete", datasource)
     assert FinetunedEmbeddingModel.open("finetuned_model_to_delete")
     FinetunedEmbeddingModel.drop("finetuned_model_to_delete")
     with pytest.raises(LookupError):

orca_sdk/memoryset.py

@@ -7,6 +7,7 @@ from typing import Any, Iterable, Literal, cast, overload
 
 import pandas as pd
 import pyarrow as pa
+from attrs import fields
 from datasets import Dataset
 from torch.utils.data import DataLoader as TorchDataLoader
 from torch.utils.data import Dataset as TorchDataset
@@ -29,11 +30,14 @@ from ._generated_api_client.api import (
     memoryset_lookup_gpu,
     potential_duplicate_groups,
     query_memoryset,
+    suggest_cascading_edits,
     update_memories_gpu,
     update_memory_gpu,
     update_memoryset,
 )
 from ._generated_api_client.models import (
+    CascadeEditSuggestionsRequest,
+    CascadingEditSuggestion,
     CloneLabeledMemorysetRequest,
     CreateLabeledMemorysetRequest,
     DeleteMemoriesRequest,
@@ -1180,6 +1184,63 @@ class LabeledMemoryset:
         updated_memories = [LabeledMemory(self.id, memory) for memory in response]
         return updated_memories[0] if isinstance(updates, dict) else updated_memories
 
+    def get_cascading_edits_suggestions(
+        self: LabeledMemoryset,
+        memory: LabeledMemory,
+        *,
+        old_label: int,
+        new_label: int,
+        max_neighbors: int = 50,
+        max_validation_neighbors: int = 10,
+        similarity_threshold: float | None = None,
+        only_if_has_old_label: bool = True,
+        exclude_if_new_label: bool = True,
+        suggestion_cooldown_time: float = 3600.0 * 24.0,  # 1 day
+        label_confirmation_cooldown_time: float = 3600.0 * 24.0 * 7,  # 1 week
+    ) -> list[CascadingEditSuggestion]:
+        """
+        Suggests cascading edits for a given memory based on nearby points with similar labels.
+
+        This function is triggered after a user changes a memory's label. It looks for nearby
+        candidates in embedding space that may be subject to similar relabeling and returns them
+        as suggestions. The system uses scoring heuristics, label filters, and cooldown tracking
+        to reduce noise and improve usability.
+
+        Params:
+            memory: The memory whose label was just changed.
+            old_label: The label this memory used to have.
+            new_label: The label it was changed to.
+            max_neighbors: Maximum number of neighbors to consider.
+            max_validation_neighbors: Maximum number of neighbors to use for label suggestion.
+            similarity_threshold: If set, only include neighbors with a lookup score above this threshold.
+            only_if_has_old_label: If True, only consider neighbors that have the old label.
+            exclude_if_new_label: If True, exclude neighbors that already have the new label.
+            suggestion_cooldown_time: Minimum time (in seconds) since the last suggestion for a neighbor
+                to be considered again.
+            label_confirmation_cooldown_time: Minimum time (in seconds) since a neighbor's label was confirmed
+                to be considered for suggestions.
+            _current_time: Optional override for the current timestamp (useful for testing).
+
+        Returns:
+            A list of CascadingEditSuggestion objects, each containing a neighbor and the suggested new label.
+        """
+
+        return suggest_cascading_edits(
+            name_or_id=self.id,
+            memory_id=memory.memory_id,
+            body=CascadeEditSuggestionsRequest(
+                old_label=old_label,
+                new_label=new_label,
+                max_neighbors=max_neighbors,
+                max_validation_neighbors=max_validation_neighbors,
+                similarity_threshold=similarity_threshold,
+                only_if_has_old_label=only_if_has_old_label,
+                exclude_if_new_label=exclude_if_new_label,
+                suggestion_cooldown_time=suggestion_cooldown_time,
+                label_confirmation_cooldown_time=label_confirmation_cooldown_time,
+            ),
+        )
+
     def delete(self, memory_id: str | Iterable[str]) -> None:
         """
         Delete memories from the memoryset
@@ -1229,6 +1290,9 @@ class LabeledMemoryset:
         Returns:
             dictionary with aggregate metrics for each analysis that was run
 
+        Raises:
+            ValueError: If an invalid analysis name is provided
+
         Examples:
             Run label and duplicate analysis:
             >>> memoryset.analyze("label", {"name": "duplicate", "possible_duplicate_threshold": 0.99})
@@ -1263,12 +1327,26 @@ class LabeledMemoryset:
             Display label analysis to review potential mislabelings:
             >>> memoryset.display_label_analysis()
         """
+
+        # Get valid analysis names from MemorysetAnalysisConfigs
+        valid_analysis_names = {
+            field.name for field in fields(MemorysetAnalysisConfigs) if field.name != "additional_properties"
+        }
+
         configs: dict[str, dict] = {}
         for analysis in analyses:
             if isinstance(analysis, str):
+                error_msg = (
+                    f"Invalid analysis name: {analysis}. Valid names are: {', '.join(sorted(valid_analysis_names))}"
+                )
+                if analysis not in valid_analysis_names:
+                    raise ValueError(error_msg)
                 configs[analysis] = {}
             else:
                 name = analysis.pop("name")  # type: ignore
+                error_msg = f"Invalid analysis name: {name}. Valid names are: {', '.join(sorted(valid_analysis_names))}"
+                if name not in valid_analysis_names:
+                    raise ValueError(error_msg)
                 configs[name] = analysis  # type: ignore
 
         analysis = analyze_memoryset(