orca-sdk 0.0.102__py3-none-any.whl → 0.0.104__py3-none-any.whl

orca_sdk/conftest.py

@@ -24,7 +24,7 @@ os.environ["ORCA_SAVE_TELEMETRY_SYNCHRONOUSLY"] = "true"
 
 def skip_in_prod(reason: str):
     """Custom decorator to skip tests when running against production API"""
-    PROD_API_URLs = ["https://api.orcadb.ai", "https://api.dev.orcadb.ai"]
+    PROD_API_URLs = ["https://api.orcadb.ai", "https://api.staging.orcadb.ai"]
     return pytest.mark.skipif(
         os.environ["ORCA_API_URL"] in PROD_API_URLs,
         reason=reason,
@@ -45,7 +45,7 @@ def _create_org_id():
 
 
 @pytest.fixture()
-def base_url_reset():
+def api_url_reset():
     original_base_url = orca_api.base_url
     yield
     orca_api.base_url = original_base_url
@@ -113,6 +113,18 @@ SAMPLE_DATA = [
     {"value": "cats have nine lives", "label": 1, "key": "g2", "score": 0.9, "source_id": "s14"},
     {"value": "tomato soup with grilled cheese", "label": 0, "key": "g1", "score": 0.1, "source_id": "s15"},
     {"value": "cats are independent animals", "label": 1, "key": "g2", "score": 0.9, "source_id": "s16"},
+    {"value": "the beach is always fun", "label": None, "key": "g3", "score": None, "source_id": "s17"},
+    {"value": "i love the beach", "label": None, "key": "g3", "score": None, "source_id": "s18"},
+    {"value": "the ocean is healing", "label": None, "key": "g3", "score": None, "source_id": "s19"},
+    {
+        "value": "sandy feet, sand between my toes at the beach",
+        "label": None,
+        "key": "g3",
+        "score": None,
+        "source_id": "s20",
+    },
+    {"value": "i am such a beach bum", "label": None, "key": "g3", "score": None, "source_id": "s21"},
+    {"value": "i will always want to be at the beach", "label": None, "key": "g3", "score": None, "source_id": "s22"},
 ]
 
 

orca_sdk/credentials.py

@@ -35,11 +35,33 @@ class OrcaCredentials:
     """
 
     @staticmethod
-    def get_api_url() -> str:
+    def is_authenticated() -> bool:
         """
-        Get the Orca API base URL that is currently being used
+        Check if you are authenticated to interact with the Orca API
+
+        Returns:
+            True if you are authenticated, False otherwise
         """
-        return str(orca_api.base_url)
+        try:
+            return orca_api.GET("/auth")
+        except ValueError as e:
+            if "Invalid API key" in str(e):
+                return False
+            raise e
+
+    @staticmethod
+    def is_healthy() -> bool:
+        """
+        Check whether the API is healthy
+
+        Returns:
+            True if the API is healthy, False otherwise
+        """
+        try:
+            orca_api.GET("/check/healthy")
+        except Exception:
+            return False
+        return True
 
     @staticmethod
     def list_api_keys() -> list[ApiKeyInfo]:
@@ -58,21 +80,6 @@ class OrcaCredentials:
             for api_key in orca_api.GET("/auth/api_key")
         ]
 
-    @staticmethod
-    def is_authenticated() -> bool:
-        """
-        Check if you are authenticated to interact with the Orca API
-
-        Returns:
-            True if you are authenticated, False otherwise
-        """
-        try:
-            return orca_api.GET("/auth")
-        except ValueError as e:
-            if "Invalid API key" in str(e):
-                return False
-            raise e
-
     @staticmethod
     def create_api_key(name: str, scopes: set[Scope] = {"ADMINISTER"}) -> str:
         """
@@ -104,20 +111,6 @@ class OrcaCredentials:
         """
         orca_api.DELETE("/auth/api_key/{name_or_id}", params={"name_or_id": name})
 
-    @staticmethod
-    def set_headers(headers: dict[str, str]):
-        """
-        Add or override default HTTP headers for all Orca API requests.
-
-        Params:
-            headers: Mapping of header names to their string values
-
-        Notes:
-            New keys are merged into the existing headers, this will overwrite headers with the
-            same name, but leave other headers untouched.
-        """
-        orca_api.headers.update(Headers(headers))
-
     @staticmethod
     def set_api_key(api_key: str, check_validity: bool = True):
         """
@@ -133,17 +126,24 @@ class OrcaCredentials:
         Raises:
             ValueError: if the API key is invalid and `check_validity` is True
         """
-        OrcaCredentials.set_headers({"Api-Key": api_key})
+        OrcaCredentials.set_api_headers({"Api-Key": api_key})
         if check_validity:
             orca_api.GET("/auth")
 
     @staticmethod
-    def set_base_url(base_url: str, check_validity: bool = True):
+    def get_api_url() -> str:
+        """
+        Get the base URL of the Orca API that is currently being used
+        """
+        return str(orca_api.base_url)
+
+    @staticmethod
+    def set_api_url(url: str, check_validity: bool = True):
         """
         Set the base URL for the Orca API
 
         Args:
-            base_url: The base URL to set
+            url: The base URL to set
             check_validity: Whether to check if there is an API running at the given base URL
 
         Raises:
@@ -152,27 +152,26 @@ class OrcaCredentials:
         # check if the base url is reachable before setting it
         if check_validity:
             try:
-                httpx.get(base_url, timeout=1)
+                httpx.get(url, timeout=1)
             except ConnectError as e:
-                raise ValueError(f"No API found at {base_url}") from e
+                raise ValueError(f"No API found at {url}") from e
 
-        orca_api.base_url = base_url
+        orca_api.base_url = url
 
         # check if the api passes the health check
         if check_validity:
-            orca_api.GET("/")
+            OrcaCredentials.is_healthy()
 
     @staticmethod
-    def is_healthy() -> bool:
+    def set_api_headers(headers: dict[str, str]):
         """
-        Check whether the API is healthy
+        Add or override default HTTP headers for all Orca API requests.
 
-        Returns:
-            True if the API is healthy, False otherwise
+        Params:
+            headers: Mapping of header names to their string values
+
+        Notes:
+            New keys are merged into the existing headers, this will overwrite headers with the
+            same name, but leave other headers untouched.
         """
-        try:
-            orca_api.GET("/")
-            orca_api.GET("/gpu/")
-        except Exception:
-            return False
-        return True
+        orca_api.headers.update(Headers(headers))

orca_sdk/credentials_test.py

@@ -38,20 +38,20 @@ def test_set_invalid_api_key(api_key):
     assert not OrcaCredentials.is_authenticated()
 
 
-def test_set_base_url(base_url_reset):
-    OrcaCredentials.set_base_url("http://api.orcadb.ai")
+def test_set_api_url(api_url_reset):
+    OrcaCredentials.set_api_url("http://api.orcadb.ai")
     assert str(orca_api.base_url) == "http://api.orcadb.ai"
 
 
 def test_set_invalid_base_url():
     with pytest.raises(ValueError, match="No API found at http://localhost:1582"):
-        OrcaCredentials.set_base_url("http://localhost:1582")
+        OrcaCredentials.set_api_url("http://localhost:1582")
 
 
 def test_is_healthy():
     assert OrcaCredentials.is_healthy()
 
 
-def test_is_healthy_false(base_url_reset):
-    OrcaCredentials.set_base_url("http://localhost:1582", check_validity=False)
+def test_is_healthy_false(api_url_reset):
+    OrcaCredentials.set_api_url("http://localhost:1582", check_validity=False)
     assert not OrcaCredentials.is_healthy()

orca_sdk/datasource.py

@@ -499,7 +499,7 @@ class Datasource:
         with open(output_path, "wb") as download_file:
             with orca_api.stream("GET", f"/datasource/{self.id}/download", params={"file_type": file_type}) as response:
                 total_chunks = int(response.headers["X-Total-Chunks"]) if "X-Total-Chunks" in response.headers else None
-                with tqdm(desc=f"Downloading", total=total_chunks, disable=total_chunks is None) as progress:
+                with tqdm(desc="Downloading", total=total_chunks, disable=total_chunks is None) as progress:
                     for chunk in response.iter_bytes():
                         download_file.write(chunk)
                         progress.update(1)

orca_sdk/datasource_test.py

@@ -329,4 +329,9 @@ def test_download_datasource(hf_dataset, datasource):
             dataset_from_downloaded_csv.remove_columns("score").to_dict()
             == hf_dataset.remove_columns("score").to_dict()
         )
-        assert np.allclose(dataset_from_downloaded_csv["score"], hf_dataset["score"])
+        # Replace None with NaN for comparison
+        assert np.allclose(
+            np.array([np.nan if v is None else float(v) for v in dataset_from_downloaded_csv["score"]], dtype=float),
+            np.array([np.nan if v is None else float(v) for v in hf_dataset["score"]], dtype=float),
+            equal_nan=True,
+        )

orca_sdk/embedding_model.py

@@ -231,7 +231,34 @@ class EmbeddingModelBase(ABC):
             else:
                 raise ValueError("Invalid embedding model")
             assert res is not None
-            return RegressionMetrics(**res) if "mse" in res else ClassificationMetrics(**res)
+            return (
+                RegressionMetrics(
+                    coverage=res.get("coverage"),
+                    mse=res.get("mse"),
+                    rmse=res.get("rmse"),
+                    mae=res.get("mae"),
+                    r2=res.get("r2"),
+                    explained_variance=res.get("explained_variance"),
+                    loss=res.get("loss"),
+                    anomaly_score_mean=res.get("anomaly_score_mean"),
+                    anomaly_score_median=res.get("anomaly_score_median"),
+                    anomaly_score_variance=res.get("anomaly_score_variance"),
+                )
+                if "mse" in res
+                else ClassificationMetrics(
+                    coverage=res.get("coverage"),
+                    f1_score=res.get("f1_score"),
+                    accuracy=res.get("accuracy"),
+                    loss=res.get("loss"),
+                    anomaly_score_mean=res.get("anomaly_score_mean"),
+                    anomaly_score_median=res.get("anomaly_score_median"),
+                    anomaly_score_variance=res.get("anomaly_score_variance"),
+                    roc_auc=res.get("roc_auc"),
+                    pr_auc=res.get("pr_auc"),
+                    pr_curve=res.get("pr_curve"),
+                    roc_curve=res.get("roc_curve"),
+                )
+            )
 
         job = Job(response["task_id"], lambda: get_result(response["task_id"]))
         return job if background else job.result()

orca_sdk/job.py

@@ -7,7 +7,7 @@ from typing import Callable, Generic, TypedDict, TypeVar, cast
 
 from tqdm.auto import tqdm
 
-from .client import TaskStatus, orca_api
+from .client import orca_api
 
 
 class JobConfig(TypedDict):
@@ -26,6 +26,9 @@ class Status(Enum):
     DISPATCHED = "DISPATCHED"
     """The job has been queued and is waiting to be processed"""
 
+    WAITING = "WAITING"
+    """The job is waiting for dependencies to complete"""
+
     PROCESSING = "PROCESSING"
     """The job is being processed"""
 

orca_sdk/job_test.py

@@ -1,10 +1,20 @@
 import time
 
+import pytest
+from datasets import Dataset
+
 from .classification_model import ClassificationModel
 from .datasource import Datasource
 from .job import Job, Status
 
 
+@pytest.fixture(scope="session")
+def datasource_without_nones(hf_dataset: Dataset):
+    return Datasource.from_hf_dataset(
+        "test_datasource_without_nones", hf_dataset.filter(lambda x: x["label"] is not None)
+    )
+
+
 def wait_for_jobs_status(job_ids, expected_statuses, timeout=10, poll_interval=0.2):
     """
     Wait until all jobs reach one of the expected statuses or timeout is reached.
@@ -18,8 +28,8 @@ def wait_for_jobs_status(job_ids, expected_statuses, timeout=10, poll_interval=0
     raise TimeoutError(f"Jobs did not reach statuses {expected_statuses} within {timeout} seconds")
 
 
-def test_job_creation(classification_model: ClassificationModel, datasource: Datasource):
-    job = classification_model.evaluate(datasource, background=True)
+def test_job_creation(classification_model: ClassificationModel, datasource_without_nones: Datasource):
+    job = classification_model.evaluate(datasource_without_nones, background=True)
     assert job.id is not None
     assert job.type == "EVALUATE_MODEL"
     assert job.status in [Status.DISPATCHED, Status.PROCESSING]
@@ -29,8 +39,8 @@ def test_job_creation(classification_model: ClassificationModel, datasource: Dat
     assert len(Job.query(limit=5, type="EVALUATE_MODEL")) >= 1
 
 
-def test_job_result(classification_model: ClassificationModel, datasource: Datasource):
-    job = classification_model.evaluate(datasource, background=True)
+def test_job_result(classification_model: ClassificationModel, datasource_without_nones: Datasource):
+    job = classification_model.evaluate(datasource_without_nones, background=True)
     result = job.result(show_progress=False)
     assert result is not None
     assert job.status == Status.COMPLETED
@@ -38,8 +48,8 @@ def test_job_result(classification_model: ClassificationModel, datasource: Datas
     assert job.steps_completed == job.steps_total
 
 
-def test_job_wait(classification_model: ClassificationModel, datasource: Datasource):
-    job = classification_model.evaluate(datasource, background=True)
+def test_job_wait(classification_model: ClassificationModel, datasource_without_nones: Datasource):
+    job = classification_model.evaluate(datasource_without_nones, background=True)
     job.wait(show_progress=False)
     assert job.status == Status.COMPLETED
     assert job.steps_completed is not None
@@ -47,8 +57,8 @@ def test_job_wait(classification_model: ClassificationModel, datasource: Datasou
     assert job.value is not None
 
 
-def test_job_refresh(classification_model: ClassificationModel, datasource: Datasource):
-    job = classification_model.evaluate(datasource, background=True)
+def test_job_refresh(classification_model: ClassificationModel, datasource_without_nones: Datasource):
+    job = classification_model.evaluate(datasource_without_nones, background=True)
     last_refreshed_at = job.refreshed_at
     # accessing the status attribute should refresh the job after the refresh interval
     Job.set_config(refresh_interval=1)
@@ -61,12 +71,12 @@ def test_job_refresh(classification_model: ClassificationModel, datasource: Data
     assert job.refreshed_at > last_refreshed_at
 
 
-def test_job_query_pagination(classification_model: ClassificationModel, datasource: Datasource):
+def test_job_query_pagination(classification_model: ClassificationModel, datasource_without_nones: Datasource):
     """Test pagination with Job.query() method"""
     # Create multiple jobs to test pagination
     jobs_created = []
     for i in range(3):
-        job = classification_model.evaluate(datasource, background=True)
+        job = classification_model.evaluate(datasource_without_nones, background=True)
         jobs_created.append(job.id)
 
     # Wait for jobs to be at least PROCESSING or COMPLETED

orca_sdk/memoryset.py

@@ -21,7 +21,9 @@ from .client import (
     FilterItem,
 )
 from .client import LabeledMemory as LabeledMemoryResponse
-from .client import LabeledMemoryInsert
+from .client import (
+    LabeledMemoryInsert,
+)
 from .client import LabeledMemoryLookup as LabeledMemoryLookupResponse
 from .client import (
     LabeledMemoryUpdate,
@@ -35,7 +37,9 @@ from .client import (
     MemoryType,
 )
 from .client import ScoredMemory as ScoredMemoryResponse
-from .client import ScoredMemoryInsert
+from .client import (
+    ScoredMemoryInsert,
+)
 from .client import ScoredMemoryLookup as ScoredMemoryLookupResponse
 from .client import (
     ScoredMemoryUpdate,
@@ -177,7 +181,7 @@ def _parse_memory_insert(memory: dict[str, Any], type: MemoryType) -> LabeledMem
     match type:
         case "LABELED":
             label = memory.get("label")
-            if not isinstance(label, int):
+            if label is not None and not isinstance(label, int):
                 raise ValueError("Memory label must be an integer")
             metadata = {k: v for k, v in memory.items() if k not in DEFAULT_COLUMN_NAMES | {"label"}}
             if any(k in metadata for k in FORBIDDEN_METADATA_COLUMN_NAMES):
@@ -187,7 +191,7 @@ def _parse_memory_insert(memory: dict[str, Any], type: MemoryType) -> LabeledMem
             return {"value": value, "label": label, "source_id": source_id, "metadata": metadata}
         case "SCORED":
             score = memory.get("score")
-            if not isinstance(score, (int, float)):
+            if score is not None and not isinstance(score, (int, float)):
                 raise ValueError("Memory score must be a number")
             metadata = {k: v for k, v in memory.items() if k not in DEFAULT_COLUMN_NAMES | {"score"}}
             if any(k in metadata for k in FORBIDDEN_METADATA_COLUMN_NAMES):
@@ -288,27 +292,13 @@ class MemoryBase(ABC):
             raise AttributeError(f"{key} is not a valid attribute")
         return self.metadata[key]
 
-    def update(
+    def _update(
         self,
         *,
         value: str = UNSET,
         source_id: str | None = UNSET,
         **metadata: None | bool | float | int | str,
     ) -> Self:
-        """
-        Update the memory with new values
-
-        Note:
-            If a field is not provided, it will default to [UNSET][orca_sdk.UNSET] and not be updated.
-
-        Params:
-            value: New value of the memory
-            source_id: New source ID of the memory
-            **metadata: New values for metadata properties
-
-        Returns:
-            The updated memory
-        """
         response = orca_api.PATCH(
             "/gpu/memoryset/{name_or_id}/memory",
             params={"name_or_id": self.memoryset_id},
@@ -368,7 +358,7 @@ class LabeledMemory(MemoryBase):
     * **`...`** (<code>[str][str] | [float][float] | [int][int] | [bool][bool] | None</code>): All metadata properties can be accessed as attributes
     """
 
-    label: int
+    label: int | None
     label_name: str | None
     memory_type = "LABELED"
 
@@ -403,7 +393,7 @@ class LabeledMemory(MemoryBase):
         self,
         *,
         value: str = UNSET,
-        label: int = UNSET,
+        label: int | None = UNSET,
         source_id: str | None = UNSET,
         **metadata: None | bool | float | int | str,
     ) -> LabeledMemory:
@@ -422,7 +412,7 @@ class LabeledMemory(MemoryBase):
         Returns:
             The updated memory
         """
-        super().update(value=value, label=label, source_id=source_id, **metadata)
+        self._update(value=value, label=label, source_id=source_id, **metadata)
         return self
 
     def to_dict(self) -> dict[str, Any]:
@@ -507,7 +497,7 @@ class ScoredMemory(MemoryBase):
     * **`...`** (<code>[str][str] | [float][float] | [int][int] | [bool][bool] | None</code>): All metadata properties can be accessed as attributes
     """
 
-    score: float
+    score: float | None
     memory_type = "SCORED"
 
     def __init__(
@@ -540,7 +530,7 @@ class ScoredMemory(MemoryBase):
         self,
         *,
         value: str = UNSET,
-        score: float = UNSET,
+        score: float | None = UNSET,
         source_id: str | None = UNSET,
         **metadata: None | bool | float | int | str,
     ) -> ScoredMemory:
@@ -559,7 +549,7 @@ class ScoredMemory(MemoryBase):
         Returns:
             The updated memory
         """
-        super().update(value=value, score=score, source_id=source_id, **metadata)
+        self._update(value=value, score=score, source_id=source_id, **metadata)
         return self
 
     def to_dict(self) -> dict[str, Any]:
@@ -645,6 +635,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
     embedding_model: EmbeddingModelBase
     index_type: IndexType
     index_params: dict[str, Any]
+    hidden: bool
 
     def __init__(self, metadata: MemorysetMetadata):
         # for internal use only, do not document
@@ -665,6 +656,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         self.index_type = metadata["index_type"]
         self.index_params = metadata["index_params"]
         self.memory_type = metadata["memory_type"]
+        self.hidden = metadata["hidden"]
 
     def __eq__(self, other) -> bool:
         return isinstance(other, MemorysetBase) and self.id == other.id
@@ -699,6 +691,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         index_params: dict[str, Any] = {},
         if_exists: CreateMode = "error",
         background: Literal[True],
+        hidden: bool = False,
     ) -> Job[Self]:
         pass
 
@@ -723,6 +716,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         index_params: dict[str, Any] = {},
         if_exists: CreateMode = "error",
         background: Literal[False] = False,
+        hidden: bool = False,
     ) -> Self:
         pass
 
@@ -746,6 +740,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         index_params: dict[str, Any] = {},
         if_exists: CreateMode = "error",
         background: bool = False,
+        hidden: bool = False,
     ) -> Self | Job[Self]:
         """
         Create a new memoryset in the OrcaCloud
@@ -783,6 +778,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             if_exists: What to do if a memoryset with the same name already exists, defaults to
                 `"error"`. Other option is `"open"` to open the existing memoryset.
             background: Whether to run the operation none blocking and return a job handle
+            hidden: Whether the memoryset should be hidden
 
         Returns:
             Handle to the new memoryset in the OrcaCloud
@@ -820,6 +816,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             "remove_duplicates": remove_duplicates,
             "index_type": index_type,
             "index_params": index_params,
+            "hidden": hidden,
         }
         if prompt is not None:
             payload["prompt"] = prompt
@@ -1272,14 +1269,20 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             return False
 
     @classmethod
-    def all(cls) -> list[Self]:
+    def all(cls, show_hidden: bool = False) -> list[Self]:
         """
         Get a list of handles to all memorysets in the OrcaCloud
 
+        Params:
+            show_hidden: Whether to include hidden memorysets in results, defaults to `False`
+
         Returns:
             List of handles to all memorysets in the OrcaCloud
         """
-        return [cls(metadata) for metadata in orca_api.GET("/memoryset", params={"type": cls.memory_type})]
+        return [
+            cls(metadata)
+            for metadata in orca_api.GET("/memoryset", params={"type": cls.memory_type, "show_hidden": show_hidden})
+        ]
 
     @classmethod
     def drop(cls, name_or_id: str, if_not_exists: DropMode = "error"):
@@ -1301,7 +1304,14 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             if if_not_exists == "error":
                 raise
 
-    def set(self, *, name: str = UNSET, description: str | None = UNSET, label_names: list[str] = UNSET):
+    def set(
+        self,
+        *,
+        name: str = UNSET,
+        description: str | None = UNSET,
+        label_names: list[str] = UNSET,
+        hidden: bool = UNSET,
+    ):
         """
         Update editable attributes of the memoryset
 
@@ -1320,6 +1330,8 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             payload["description"] = description
         if label_names is not UNSET:
             payload["label_names"] = label_names
+        if hidden is not UNSET:
+            payload["hidden"] = hidden
 
         orca_api.PATCH("/memoryset/{name_or_id}", params={"name_or_id": self.id}, json=payload)
         self.refresh()

orca_sdk/memoryset_test.py

@@ -122,6 +122,26 @@ def test_all_memorysets(readonly_memoryset: LabeledMemoryset):
     assert any(memoryset.name == readonly_memoryset.name for memoryset in memorysets)
 
 
+def test_all_memorysets_hidden(
+    readonly_memoryset: LabeledMemoryset,
+):
+    # Create a hidden memoryset
+    hidden_memoryset = LabeledMemoryset.clone(readonly_memoryset, "test_hidden_memoryset")
+    hidden_memoryset.set(hidden=True)
+
+    # Test that show_hidden=False excludes hidden memorysets
+    visible_memorysets = LabeledMemoryset.all(show_hidden=False)
+    assert len(visible_memorysets) > 0
+    assert readonly_memoryset in visible_memorysets
+    assert hidden_memoryset not in visible_memorysets
+
+    # Test that show_hidden=True includes hidden memorysets
+    all_memorysets = LabeledMemoryset.all(show_hidden=True)
+    assert len(all_memorysets) == len(visible_memorysets) + 1
+    assert readonly_memoryset in all_memorysets
+    assert hidden_memoryset in all_memorysets
+
+
 def test_all_memorysets_unauthenticated(unauthenticated):
     with pytest.raises(ValueError, match="Invalid API key"):
         LabeledMemoryset.all()
@@ -167,6 +187,9 @@ def test_update_memoryset_attributes(writable_memoryset: LabeledMemoryset):
     writable_memoryset.set(label_names=["New label 1", "New label 2"])
     assert writable_memoryset.label_names == ["New label 1", "New label 2"]
 
+    writable_memoryset.set(hidden=True)
+    assert writable_memoryset.hidden is True
+
 
 def test_search(readonly_memoryset: LabeledMemoryset):
     memory_lookups = readonly_memoryset.search(["i love soup", "cats are cute"])
@@ -364,7 +387,7 @@ def test_clone_memoryset(readonly_memoryset: LabeledMemoryset):
 
 def test_embedding_evaluation(eval_datasource: Datasource):
     results = LabeledMemoryset.run_embedding_evaluation(
-        eval_datasource, embedding_models=["CDE_SMALL"], neighbor_count=2
+        eval_datasource, embedding_models=["CDE_SMALL"], neighbor_count=3
     )
     assert isinstance(results, list)
     assert len(results) == 1
@@ -465,13 +488,14 @@ def test_drop_memoryset(writable_memoryset: LabeledMemoryset):
 
 
 def test_scored_memoryset(scored_memoryset: ScoredMemoryset):
-    assert scored_memoryset.length == 16
+    assert scored_memoryset.length == 22
     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 is not None
     assert lookup[0].score < 0.11