orca-sdk 0.1.11__py3-none-any.whl → 0.1.12__py3-none-any.whl

orca_sdk/memoryset.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import logging
 from abc import ABC
 from datetime import datetime, timedelta
 from os import PathLike
@@ -16,9 +15,7 @@ from typing import (
     overload,
 )
 
-from datasets import Dataset
-
-from ._utils.common import UNSET, CreateMode, DropMode
+from ._utils.common import UNSET, CreateMode, DropMode, logger
 from .async_client import OrcaAsyncClient
 from .client import (
     CascadingEditSuggestion,
@@ -73,6 +70,7 @@ from .telemetry import ClassificationPrediction, RegressionPrediction
 
 if TYPE_CHECKING:
     # peer dependencies that are used for types only
+    from datasets import Dataset as HFDataset  # type: ignore
     from pandas import DataFrame as PandasDataFrame  # type: ignore
     from pyarrow import Table as PyArrowTable  # type: ignore
     from torch.utils.data import DataLoader as TorchDataLoader  # type: ignore
@@ -98,6 +96,16 @@ FilterOperation = Literal["==", "!=", ">", ">=", "<", "<=", "in", "not in", "lik
 Operations that can be used in a filter expression.
 """
 
+ConsistencyLevel = Literal["Strong", "Session", "Bounded", "Eventual"]
+"""
+Consistency level for memoryset reads.
+
+* **`Strong`**: Reads include all committed writes; may wait for full freshness.
+* **`Session`**: Reads include all writes that happened in the same server process.
+* **`Bounded`**: Reads may miss newest writes within a small staleness window.
+* **`Eventual`**: No freshness guarantee; reads can miss recent writes.
+"""
+
 FilterValue = (
     str
     | int
@@ -989,6 +997,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         description: Description of the memoryset
         length: Number of memories in the memoryset
         embedding_model: Embedding model used to embed the memory values for semantic search
+        partitioned: Whether the memoryset is partitioned
         created_at: When the memoryset was created, automatically generated on create
         updated_at: When the memoryset was last updated, automatically updated on updates
     """
@@ -1003,6 +1012,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
     updated_at: datetime
     insertion_status: Status | None
     embedding_model: EmbeddingModelBase
+    partitioned: bool
     index_type: IndexType
     index_params: dict[str, Any]
     hidden: bool
@@ -1024,6 +1034,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         self.insertion_status = (
             Status(metadata["insertion_status"]) if metadata["insertion_status"] is not None else None
         )
+        self.partitioned = metadata["is_partitioned"]
         self._last_refresh = datetime.now()
         self.index_type = metadata["index_type"]
         self.index_params = metadata["index_params"]
@@ -1085,6 +1096,60 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
 
         return existing
 
+    @classmethod
+    def _create_empty(
+        cls,
+        name: str,
+        *,
+        embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
+        partitioned: bool = False,
+        description: str | None = None,
+        label_names: list[str] | None = None,
+        max_seq_length_override: int | None = None,
+        prompt: str | None = None,
+        index_type: IndexType = "FLAT",
+        index_params: dict[str, Any] = {},
+        if_exists: CreateMode = "error",
+        hidden: bool = False,
+        memory_type: MemoryType | None = None,
+    ) -> Self:
+        if embedding_model is None:
+            embedding_model = PretrainedEmbeddingModel.GTE_BASE
+
+        existing = cls._handle_if_exists(
+            name,
+            if_exists=if_exists,
+            label_names=label_names,
+            embedding_model=embedding_model,
+        )
+        if existing is not None:
+            return existing
+
+        payload: CreateMemorysetRequest = {
+            "name": name,
+            "description": description,
+            "label_names": label_names,
+            "max_seq_length_override": max_seq_length_override,
+            "index_type": index_type,
+            "index_params": index_params,
+            "hidden": hidden,
+            "is_partitioned": partitioned,
+        }
+        if memory_type is not None:
+            payload["memory_type"] = memory_type
+        if prompt is not None:
+            payload["prompt"] = prompt
+        if isinstance(embedding_model, PretrainedEmbeddingModel):
+            payload["pretrained_embedding_model_name"] = embedding_model.name
+        elif isinstance(embedding_model, FinetunedEmbeddingModel):
+            payload["finetuned_embedding_model_name_or_id"] = embedding_model.id
+        else:
+            raise ValueError("Invalid embedding model")
+
+        client = OrcaClient._resolve_client()
+        response = client.POST("/memoryset/empty", json=payload)
+        return cls.open(response["id"])
+
     @classmethod
     def _create_from_datasource(
         cls,
@@ -1110,11 +1175,10 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         subsample: int | float | None = None,
         memory_type: MemoryType | None = None,
     ) -> Self | Job[Self]:
-        """
-        Create a memoryset from a datasource by calling the API.
-
-        This is a private method that performs the actual API call to create a memoryset from a datasource.
-        """
+        # Private method for the actual API call to create a memoryset from a datasource.
+        # This exists because subclass from_datasource() methods have narrower signatures
+        # (e.g., ScoredMemoryset only has score_column, not label_column), so they can't
+        # be called polymorphically. Both create() and from_datasource() delegate here.
         if embedding_model is None:
             embedding_model = PretrainedEmbeddingModel.GTE_BASE
 
@@ -1172,6 +1236,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         *,
         datasource: None = None,
         embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
+        partitioned: bool = False,
         description: str | None = None,
         label_names: list[str] | None = None,
         max_seq_length_override: int | None = None,
@@ -1240,6 +1305,35 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
     ) -> Self:
         pass
 
+    @overload
+    @classmethod
+    def create(
+        cls,
+        name: str,
+        *,
+        datasource: Datasource | None = None,
+        embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
+        value_column: str = "value",
+        label_column: str | None = None,
+        score_column: str | None = None,
+        source_id_column: str | None = None,
+        partition_id_column: str | None = None,
+        partitioned: bool = False,
+        description: str | None = None,
+        label_names: list[str] | None = None,
+        max_seq_length_override: int | None = None,
+        prompt: str | None = None,
+        remove_duplicates: bool = True,
+        index_type: IndexType = "FLAT",
+        index_params: dict[str, Any] = {},
+        if_exists: CreateMode = "error",
+        background: bool = False,
+        hidden: bool = False,
+        subsample: int | float | None = None,
+        memory_type: MemoryType | None = None,
+    ) -> Self | Job[Self]:
+        pass
+
     @classmethod
     def create(
         cls,
@@ -1252,6 +1346,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         score_column: str | None = None,
         source_id_column: str | None = None,
         partition_id_column: str | None = None,
+        partitioned: bool = False,
         description: str | None = None,
         label_names: list[str] | None = None,
         max_seq_length_override: int | None = None,
@@ -1289,6 +1384,8 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             source_id_column: Optional name of the column in the datasource that contains the ids in
                 the system of reference
             partition_id_column: Optional name of the column in the datasource that contains the partition ids
+            partitioned: Whether the memoryset should be partitioned. Only valid when creating an
+                empty memoryset (datasource is None). Use partition_id_column when creating from a datasource.
             description: Optional description for the memoryset, this will be used in agentic flows,
                 so make sure it is concise and describes the contents of your memoryset not the
                 datasource or the embedding model.
@@ -1333,9 +1430,12 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 index_params=index_params,
                 if_exists=if_exists,
                 hidden=hidden,
+                partitioned=partitioned,
                 memory_type=memory_type,
             )
         else:
+            if partitioned:
+                raise ValueError("Use 'partition_id_column' instead of 'partitioned' when creating from a datasource")
             return cls._create_from_datasource(
                 name,
                 datasource=datasource,
@@ -1515,105 +1615,21 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             memory_type=memory_type,
         )
 
-    @classmethod
-    def _create_empty(
-        cls,
-        name: str,
-        *,
-        embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
-        description: str | None = None,
-        label_names: list[str] | None = None,
-        max_seq_length_override: int | None = None,
-        prompt: str | None = None,
-        index_type: IndexType = "FLAT",
-        index_params: dict[str, Any] = {},
-        if_exists: CreateMode = "error",
-        hidden: bool = False,
-        memory_type: MemoryType | None = None,
-    ) -> Self:
-        """
-        Create an empty memoryset in the OrcaCloud
-
-        This creates a memoryset with no initial memories. You can add memories later using
-        the `insert` method.
-
-        Params:
-            name: Name for the new memoryset (must be unique)
-            embedding_model: Embedding model to use for embedding memory values for semantic search.
-                If not provided, a default embedding model for the memoryset will be used.
-            description: Optional description for the memoryset, this will be used in agentic flows,
-                so make sure it is concise and describes the contents of your memoryset not the
-                datasource or the embedding model.
-            label_names: List of human-readable names for the labels in the memoryset
-            max_seq_length_override: Maximum sequence length of values in the memoryset, if the
-                value is longer than this it will be truncated, will default to the model's max
-                sequence length if not provided
-            prompt: Optional prompt to use when embedding documents/memories for storage
-            index_type: Type of vector index to use for the memoryset, defaults to `"FLAT"`. Valid
-                values are `"FLAT"`, `"IVF_FLAT"`, `"IVF_SQ8"`, `"IVF_PQ"`, `"HNSW"`, and `"DISKANN"`.
-            index_params: Parameters for the vector index, defaults to `{}`
-            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.
-            hidden: Whether the memoryset should be hidden
-            memory_type: Type of memoryset to create, defaults to `"LABELED"` if called from
-                `LabeledMemoryset` and `"SCORED"` if called from `ScoredMemoryset`.
-
-        Returns:
-            Handle to the new memoryset in the OrcaCloud
-
-        Raises:
-            ValueError: If the memoryset already exists and if_exists is `"error"` or if it is
-                `"open"` and the params do not match those of the existing memoryset.
-        """
-        if embedding_model is None:
-            embedding_model = PretrainedEmbeddingModel.GTE_BASE
-
-        existing = cls._handle_if_exists(
-            name,
-            if_exists=if_exists,
-            label_names=label_names,
-            embedding_model=embedding_model,
-        )
-        if existing is not None:
-            return existing
-
-        payload: CreateMemorysetRequest = {
-            "name": name,
-            "description": description,
-            "label_names": label_names,
-            "max_seq_length_override": max_seq_length_override,
-            "index_type": index_type,
-            "index_params": index_params,
-            "hidden": hidden,
-        }
-        if memory_type is not None:
-            payload["memory_type"] = memory_type
-        if prompt is not None:
-            payload["prompt"] = prompt
-        if isinstance(embedding_model, PretrainedEmbeddingModel):
-            payload["pretrained_embedding_model_name"] = embedding_model.name
-        elif isinstance(embedding_model, FinetunedEmbeddingModel):
-            payload["finetuned_embedding_model_name_or_id"] = embedding_model.id
-        else:
-            raise ValueError("Invalid embedding model")
-
-        client = OrcaClient._resolve_client()
-        response = client.POST("/memoryset/empty", json=payload)
-        return cls.open(response["id"])
-
     @overload
     @classmethod
-    def from_hf_dataset(cls, name: str, hf_dataset: Dataset, background: Literal[True], **kwargs: Any) -> Self:
+    def from_hf_dataset(cls, name: str, hf_dataset: HFDataset, background: Literal[True], **kwargs: Any) -> Self:
         pass
 
     @overload
     @classmethod
-    def from_hf_dataset(cls, name: str, hf_dataset: Dataset, background: Literal[False] = False, **kwargs: Any) -> Self:
+    def from_hf_dataset(
+        cls, name: str, hf_dataset: HFDataset, background: Literal[False] = False, **kwargs: Any
+    ) -> Self:
         pass
 
     @classmethod
     def from_hf_dataset(
-        cls, name: str, hf_dataset: Dataset, background: bool = False, **kwargs: Any
+        cls, name: str, hf_dataset: HFDataset, background: bool = False, **kwargs: Any
     ) -> Self | Job[Self]:
         """
         Create a new memoryset from a Hugging Face [`Dataset`][datasets.Dataset] in the OrcaCloud
@@ -2164,7 +2180,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         try:
             client = OrcaClient._resolve_client()
             client.DELETE("/memoryset/{name_or_id}", params={"name_or_id": name_or_id, "cascade": cascade})
-            logging.info(f"Deleted memoryset {name_or_id}")
+            logger.info(f"Deleted memoryset {name_or_id}")
         except LookupError:
             if if_not_exists == "error":
                 raise
@@ -2210,6 +2226,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         embedding_model: PretrainedEmbeddingModel | FinetunedEmbeddingModel | None = None,
         max_seq_length_override: int | None = None,
         prompt: str | None = None,
+        partitioned: bool | None = None,
         if_exists: CreateMode = "error",
         background: Literal[True],
     ) -> Job[Self]:
@@ -2223,6 +2240,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         embedding_model: PretrainedEmbeddingModel | FinetunedEmbeddingModel | None = None,
         max_seq_length_override: int | None = None,
         prompt: str | None = None,
+        partitioned: bool | None = None,
         if_exists: CreateMode = "error",
         background: Literal[False] = False,
     ) -> Self:
@@ -2235,6 +2253,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         embedding_model: PretrainedEmbeddingModel | FinetunedEmbeddingModel | None = None,
         max_seq_length_override: int | None = UNSET,
         prompt: str | None = None,
+        partitioned: bool | None = None,
         if_exists: CreateMode = "error",
         background: bool = False,
     ) -> Self | Job[Self]:
@@ -2250,6 +2269,8 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 If not provided, will use the source memoryset's max sequence length.
             prompt: Optional custom prompt to use for the cloned memoryset.
                 If not provided, will use the source memoryset's prompt.
+            partitioned: Whether the cloned memoryset should be partitioned.
+                If not provided, will inherit the source memoryset's partitioning.
             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.
 
@@ -2290,6 +2311,8 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             payload["pretrained_embedding_model_name"] = embedding_model.name
         elif isinstance(embedding_model, FinetunedEmbeddingModel):
             payload["finetuned_embedding_model_name_or_id"] = embedding_model.id
+        if partitioned is not None:
+            payload["is_partitioned"] = partitioned
 
         client = OrcaClient._resolve_client()
         metadata = client.POST("/memoryset/{name_or_id}/clone", params={"name_or_id": self.id}, json=payload)
@@ -2387,6 +2410,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         partition_filter_mode: Literal[
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
+        consistency_level: ConsistencyLevel = "Bounded",
     ) -> list[MemoryLookupT]:
         pass
 
@@ -2401,6 +2425,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         partition_filter_mode: Literal[
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
+        consistency_level: ConsistencyLevel = "Bounded",
     ) -> list[list[MemoryLookupT]]:
         pass
 
@@ -2414,6 +2439,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         partition_filter_mode: Literal[
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
+        consistency_level: ConsistencyLevel = "Bounded",
     ) -> list[MemoryLookupT] | list[list[MemoryLookupT]]:
         """
         Search for memories that are semantically similar to the query
@@ -2429,6 +2455,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 - "include_global": Include global memories
                 - "exclude_global": Exclude global memories
                 - "only_global": Only include global memories
+            consistency_level: Consistency level to use for the search
         Returns:
             List of memories from the memoryset that match the query. If a single query is provided,
                 the return value is a list containing a single list of memories. If a list of
@@ -2470,6 +2497,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 "prompt": prompt,
                 "partition_id": partition_id,
                 "partition_filter_mode": partition_filter_mode,
+                "consistency_level": consistency_level,
             },
         )
         lookups = [
@@ -2495,6 +2523,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         filters: list[FilterItemTuple] = [],
         with_feedback_metrics: bool = False,
         sort: list[TelemetrySortItem] | None = None,
+        consistency_level: ConsistencyLevel = "Bounded",
     ) -> list[MemoryT]:
         """
         Query the memoryset for memories that match the filters
@@ -2502,8 +2531,10 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         Params:
             offset: The offset of the first memory to return
             limit: The maximum number of memories to return
-            filters: List of filters to apply to the query.
+            filters: List of filters to apply to the query
             with_feedback_metrics: Whether to include feedback metrics in the response
+            sort: Optional sort order to apply
+            consistency_level: Consistency level to use for the query
 
         Returns:
             List of memories from the memoryset that match the filters
@@ -2526,6 +2557,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                     "limit": limit,
                     "filters": [_parse_filter_item_from_tuple(filter) for filter in filters],
                     "sort": [_parse_sort_item_from_tuple(item) for item in sort] if sort else None,
+                    "consistency_level": consistency_level,
                 },
             )
             return [
@@ -2540,7 +2572,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             raise ValueError("Feedback metrics are only supported when the with_feedback_metrics flag is set to True")
 
         if sort:
-            logging.warning("Sorting is not supported when with_feedback_metrics is False. Sort value will be ignored.")
+            logger.warning("Sorting is not supported when with_feedback_metrics is False. Sort value will be ignored.")
 
         response = client.POST(
             "/memoryset/{name_or_id}/memories",
@@ -2549,6 +2581,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 "offset": offset,
                 "limit": limit,
                 "filters": [_parse_filter_item_from_tuple(filter, allow_metric_fields=False) for filter in filters],
+                "consistency_level": consistency_level,
             },
         )
         return [
@@ -2686,19 +2719,22 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         self._last_refresh = current_time
 
     @overload
-    def get(self, memory_id: str) -> MemoryT:  # type: ignore -- this takes precedence
+    def get(self, memory_id: str, consistency_level: ConsistencyLevel = "Bounded") -> MemoryT:  # type: ignore -- this takes precedence
         pass
 
     @overload
-    def get(self, memory_id: Iterable[str]) -> list[MemoryT]:
+    def get(self, memory_id: Iterable[str], consistency_level: ConsistencyLevel = "Bounded") -> list[MemoryT]:
         pass
 
-    def get(self, memory_id: str | Iterable[str]) -> MemoryT | list[MemoryT]:
+    def get(
+        self, memory_id: str | Iterable[str], consistency_level: ConsistencyLevel = "Bounded"
+    ) -> MemoryT | list[MemoryT]:
         """
         Fetch a memory or memories from the memoryset
 
         Params:
             memory_id: Unique identifier of the memory or memories to fetch
+            consistency_level: Consistency level to use for the get operation
 
         Returns:
             Memory or list of memories from the memoryset
@@ -2724,7 +2760,8 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         if isinstance(memory_id, str):
             client = OrcaClient._resolve_client()
             response = client.GET(
-                "/memoryset/{name_or_id}/memory/{memory_id}", params={"name_or_id": self.id, "memory_id": memory_id}
+                "/memoryset/{name_or_id}/memory/{memory_id}",
+                params={"name_or_id": self.id, "memory_id": memory_id, "consistency_level": consistency_level},
             )
             return cast(
                 MemoryT,
@@ -2735,7 +2772,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             response = client.POST(
                 "/memoryset/{name_or_id}/memories/get",
                 params={"name_or_id": self.id},
-                json={"memory_ids": list(memory_id)},
+                json={"memory_ids": list(memory_id), "consistency_level": consistency_level},
             )
             return [
                 cast(
@@ -2982,7 +3019,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 )
                 total_deleted += response.get("deleted_count", 0)
             if total_deleted > 0:
-                logging.info(f"Deleted {total_deleted} memories from memoryset.")
+                logger.info(f"Deleted {total_deleted} memories from memoryset.")
                 self.refresh()
             return total_deleted
 
@@ -3000,7 +3037,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             },
         )
         deleted_count = response["deleted_count"]
-        logging.info(f"Deleted {deleted_count} memories from memoryset.")
+        logger.info(f"Deleted {deleted_count} memories from memoryset.")
         if deleted_count > 0:
             self.refresh()
         return deleted_count
@@ -3028,7 +3065,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             },
         )
         deleted_count = response["deleted_count"]
-        logging.info(f"Deleted {deleted_count} memories from memoryset.")
+        logger.info(f"Deleted {deleted_count} memories from memoryset.")
         if deleted_count > 0:
             self.refresh()
         return deleted_count
@@ -3233,6 +3270,7 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         *,
         datasource: None = None,
         embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
+        partitioned: bool = False,
         description: str | None = None,
         label_names: list[str],
         max_seq_length_override: int | None = None,
@@ -3307,6 +3345,7 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         label_column: str | None = "label",
         source_id_column: str | None = None,
         partition_id_column: str | None = None,
+        partitioned: bool = False,
         description: str | None = None,
         label_names: list[str] | None = None,
         max_seq_length_override: int | None = None,
@@ -3371,70 +3410,29 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
             ValueError: If the memoryset already exists and if_exists is `"error"` or if it is
                 `"open"` and the params do not match those of the existing memoryset.
         """
-        if datasource is None:
-            return super().create(
-                name,
-                datasource=None,
-                embedding_model=embedding_model,
-                description=description,
-                label_names=label_names,
-                max_seq_length_override=max_seq_length_override,
-                prompt=prompt,
-                index_type=index_type,
-                index_params=index_params,
-                if_exists=if_exists,
-                hidden=hidden,
-                memory_type="LABELED",
-            )
-        else:
-            # Type narrowing: datasource is definitely Datasource here
-            assert datasource is not None
-            if background:
-                return super().create(
-                    name,
-                    datasource=datasource,
-                    label_column=label_column,
-                    score_column=None,
-                    embedding_model=embedding_model,
-                    value_column=value_column,
-                    source_id_column=source_id_column,
-                    partition_id_column=partition_id_column,
-                    description=description,
-                    label_names=label_names,
-                    max_seq_length_override=max_seq_length_override,
-                    prompt=prompt,
-                    remove_duplicates=remove_duplicates,
-                    index_type=index_type,
-                    index_params=index_params,
-                    if_exists=if_exists,
-                    background=True,
-                    hidden=hidden,
-                    subsample=subsample,
-                    memory_type="LABELED",
-                )
-            else:
-                return super().create(
-                    name,
-                    datasource=datasource,
-                    label_column=label_column,
-                    score_column=None,
-                    embedding_model=embedding_model,
-                    value_column=value_column,
-                    source_id_column=source_id_column,
-                    partition_id_column=partition_id_column,
-                    description=description,
-                    label_names=label_names,
-                    max_seq_length_override=max_seq_length_override,
-                    prompt=prompt,
-                    remove_duplicates=remove_duplicates,
-                    index_type=index_type,
-                    index_params=index_params,
-                    if_exists=if_exists,
-                    background=False,
-                    hidden=hidden,
-                    subsample=subsample,
-                    memory_type="LABELED",
-                )
+        return super().create(
+            name,
+            datasource=datasource,
+            embedding_model=embedding_model,
+            value_column=value_column,
+            label_column=label_column,
+            score_column=None,
+            source_id_column=source_id_column,
+            partition_id_column=partition_id_column,
+            partitioned=partitioned,
+            description=description,
+            label_names=label_names,
+            max_seq_length_override=max_seq_length_override,
+            prompt=prompt,
+            remove_duplicates=remove_duplicates,
+            index_type=index_type,
+            index_params=index_params,
+            if_exists=if_exists,
+            background=background,
+            hidden=hidden,
+            subsample=subsample,
+            memory_type="LABELED",
+        )
 
     @overload
     @classmethod
@@ -3562,52 +3560,27 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
             ValueError: If the memoryset already exists and if_exists is `"error"` or if it is
                 `"open"` and the params do not match those of the existing memoryset.
         """
-        if background:
-            return super().create(
-                name,
-                datasource=datasource,
-                label_column=label_column,
-                score_column=None,
-                embedding_model=embedding_model,
-                value_column=value_column,
-                source_id_column=source_id_column,
-                partition_id_column=partition_id_column,
-                description=description,
-                label_names=label_names,
-                max_seq_length_override=max_seq_length_override,
-                prompt=prompt,
-                remove_duplicates=remove_duplicates,
-                index_type=index_type,
-                index_params=index_params,
-                if_exists=if_exists,
-                background=True,
-                hidden=hidden,
-                subsample=subsample,
-                memory_type="LABELED",
-            )
-        else:
-            return super().create(
-                name,
-                datasource=datasource,
-                label_column=label_column,
-                score_column=None,
-                embedding_model=embedding_model,
-                value_column=value_column,
-                source_id_column=source_id_column,
-                partition_id_column=partition_id_column,
-                description=description,
-                label_names=label_names,
-                max_seq_length_override=max_seq_length_override,
-                prompt=prompt,
-                remove_duplicates=remove_duplicates,
-                index_type=index_type,
-                index_params=index_params,
-                if_exists=if_exists,
-                background=False,
-                hidden=hidden,
-                subsample=subsample,
-                memory_type="LABELED",
-            )
+        return cls._create_from_datasource(
+            name,
+            datasource=datasource,
+            label_column=label_column,
+            embedding_model=embedding_model,
+            value_column=value_column,
+            source_id_column=source_id_column,
+            partition_id_column=partition_id_column,
+            description=description,
+            label_names=label_names,
+            max_seq_length_override=max_seq_length_override,
+            prompt=prompt,
+            remove_duplicates=remove_duplicates,
+            index_type=index_type,
+            index_params=index_params,
+            if_exists=if_exists,
+            background=background,
+            hidden=hidden,
+            subsample=subsample,
+            memory_type="LABELED",
+        )
 
     def display_label_analysis(self):
         """
@@ -3664,6 +3637,7 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         *,
         datasource: None = None,
         embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
+        partitioned: bool = False,
         description: str | None = None,
         max_seq_length_override: int | None = None,
         prompt: str | None = None,
@@ -3707,8 +3681,8 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         *,
         datasource: Datasource,
         embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
-        score_column: str | None = "score",
         value_column: str = "value",
+        score_column: str | None = "score",
         source_id_column: str | None = None,
         partition_id_column: str | None = None,
         description: str | None = None,
@@ -3735,6 +3709,7 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         score_column: str | None = "score",
         source_id_column: str | None = None,
         partition_id_column: str | None = None,
+        partitioned: bool = False,
         description: str | None = None,
         max_seq_length_override: int | None = None,
         prompt: str | None = None,
@@ -3792,65 +3767,27 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
             ValueError: If the memoryset already exists and if_exists is `"error"` or if it is
                 `"open"` and the params do not match those of the existing memoryset.
         """
-        if datasource is None:
-            return super().create(
-                name,
-                datasource=None,
-                embedding_model=embedding_model,
-                description=description,
-                max_seq_length_override=max_seq_length_override,
-                prompt=prompt,
-                index_type=index_type,
-                index_params=index_params,
-                if_exists=if_exists,
-                hidden=hidden,
-                memory_type="SCORED",
-            )
-        else:
-            # Type narrowing: datasource is definitely Datasource here
-            assert datasource is not None
-            if background:
-                return super().create(
-                    name,
-                    datasource=datasource,
-                    embedding_model=embedding_model,
-                    value_column=value_column,
-                    score_column=score_column,
-                    source_id_column=source_id_column,
-                    partition_id_column=partition_id_column,
-                    description=description,
-                    max_seq_length_override=max_seq_length_override,
-                    prompt=prompt,
-                    remove_duplicates=remove_duplicates,
-                    index_type=index_type,
-                    index_params=index_params,
-                    if_exists=if_exists,
-                    background=True,
-                    hidden=hidden,
-                    subsample=subsample,
-                    memory_type="SCORED",
-                )
-            else:
-                return super().create(
-                    name,
-                    datasource=datasource,
-                    embedding_model=embedding_model,
-                    value_column=value_column,
-                    score_column=score_column,
-                    source_id_column=source_id_column,
-                    partition_id_column=partition_id_column,
-                    description=description,
-                    max_seq_length_override=max_seq_length_override,
-                    prompt=prompt,
-                    remove_duplicates=remove_duplicates,
-                    index_type=index_type,
-                    index_params=index_params,
-                    if_exists=if_exists,
-                    background=False,
-                    hidden=hidden,
-                    subsample=subsample,
-                    memory_type="SCORED",
-                )
+        return super().create(
+            name,
+            datasource=datasource,
+            embedding_model=embedding_model,
+            value_column=value_column,
+            score_column=score_column,
+            source_id_column=source_id_column,
+            partition_id_column=partition_id_column,
+            partitioned=partitioned,
+            description=description,
+            max_seq_length_override=max_seq_length_override,
+            prompt=prompt,
+            remove_duplicates=remove_duplicates,
+            index_type=index_type,
+            index_params=index_params,
+            if_exists=if_exists,
+            background=background,
+            hidden=hidden,
+            subsample=subsample,
+            memory_type="SCORED",
+        )
 
     @overload
     @classmethod
@@ -3969,48 +3906,26 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
             ValueError: If the memoryset already exists and if_exists is `"error"` or if it is
                 `"open"` and the params do not match those of the existing memoryset.
         """
-        if background:
-            return super().create(
-                name,
-                datasource=datasource,
-                embedding_model=embedding_model,
-                value_column=value_column,
-                score_column=score_column,
-                source_id_column=source_id_column,
-                partition_id_column=partition_id_column,
-                description=description,
-                max_seq_length_override=max_seq_length_override,
-                prompt=prompt,
-                remove_duplicates=remove_duplicates,
-                index_type=index_type,
-                index_params=index_params,
-                if_exists=if_exists,
-                background=True,
-                hidden=hidden,
-                subsample=subsample,
-                memory_type="SCORED",
-            )
-        else:
-            return super().create(
-                name,
-                datasource=datasource,
-                embedding_model=embedding_model,
-                value_column=value_column,
-                score_column=score_column,
-                source_id_column=source_id_column,
-                partition_id_column=partition_id_column,
-                description=description,
-                max_seq_length_override=max_seq_length_override,
-                prompt=prompt,
-                remove_duplicates=remove_duplicates,
-                index_type=index_type,
-                index_params=index_params,
-                if_exists=if_exists,
-                background=False,
-                hidden=hidden,
-                subsample=subsample,
-                memory_type="SCORED",
-            )
+        return cls._create_from_datasource(
+            name,
+            datasource=datasource,
+            embedding_model=embedding_model,
+            value_column=value_column,
+            score_column=score_column,
+            source_id_column=source_id_column,
+            partition_id_column=partition_id_column,
+            description=description,
+            max_seq_length_override=max_seq_length_override,
+            prompt=prompt,
+            remove_duplicates=remove_duplicates,
+            index_type=index_type,
+            index_params=index_params,
+            if_exists=if_exists,
+            background=background,
+            hidden=hidden,
+            subsample=subsample,
+            memory_type="SCORED",
+        )
 
     @property
     def regression_models(self) -> list[RegressionModel]: