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

orca_sdk/memoryset.py

@@ -4,7 +4,17 @@ import logging
 from abc import ABC
 from datetime import datetime, timedelta
 from os import PathLike
-from typing import Any, Generic, Iterable, Literal, Self, TypeVar, cast, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generic,
+    Iterable,
+    Literal,
+    Self,
+    TypeVar,
+    cast,
+    overload,
+)
 
 import pandas as pd
 import pyarrow as pa
@@ -29,6 +39,7 @@ from .client import (
     LabeledMemoryUpdate,
     LabeledMemoryWithFeedbackMetrics,
     LabelPredictionMemoryLookup,
+    LabelPredictionWithMemoriesAndFeedback,
     MemoryMetrics,
     MemorysetAnalysisConfigs,
     MemorysetMetadata,
@@ -36,6 +47,7 @@ from .client import (
     MemorysetUpdate,
     MemoryType,
     OrcaClient,
+    PredictionFeedback,
 )
 from .client import ScoredMemory as ScoredMemoryResponse
 from .client import (
@@ -46,6 +58,7 @@ from .client import (
     ScoredMemoryUpdate,
     ScoredMemoryWithFeedbackMetrics,
     ScorePredictionMemoryLookup,
+    ScorePredictionWithMemoriesAndFeedback,
     TelemetryFilterItem,
     TelemetrySortOptions,
 )
@@ -56,6 +69,11 @@ from .embedding_model import (
     PretrainedEmbeddingModel,
 )
 from .job import Job, Status
+from .telemetry import ClassificationPrediction, RegressionPrediction
+
+if TYPE_CHECKING:
+    from .classification_model import ClassificationModel
+    from .regression_model import RegressionModel
 
 TelemetrySortItem = tuple[str, Literal["asc", "desc"]]
 """
@@ -74,7 +92,7 @@ FilterOperation = Literal["==", "!=", ">", ">=", "<", "<=", "in", "not in", "lik
 Operations that can be used in a filter expression.
 """
 
-FilterValue = str | int | float | bool | datetime | None | list[str] | list[int] | list[float] | list[bool]
+FilterValue = str | int | float | bool | datetime | None | list[str | None] | list[int] | list[float] | list[bool]
 """
 Values that can be used in a filter expression.
 """
@@ -96,7 +114,7 @@ Examples:
 
 IndexType = Literal["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ", "HNSW", "DISKANN"]
 
-DEFAULT_COLUMN_NAMES = {"value", "source_id"}
+DEFAULT_COLUMN_NAMES = {"value", "source_id", "partition_id"}
 TYPE_SPECIFIC_COLUMN_NAMES = {"label", "score"}
 FORBIDDEN_METADATA_COLUMN_NAMES = {
     "memory_id",
@@ -176,8 +194,11 @@ def _parse_memory_insert(memory: dict[str, Any], type: MemoryType) -> LabeledMem
     if not isinstance(value, str):
         raise ValueError("Memory value must be a string")
     source_id = memory.get("source_id")
-    if source_id and not isinstance(source_id, str):
+    if source_id is not None and not isinstance(source_id, str):
         raise ValueError("Memory source_id must be a string")
+    partition_id = memory.get("partition_id")
+    if partition_id is not None and not isinstance(partition_id, str):
+        raise ValueError("Memory partition_id must be a string")
     match type:
         case "LABELED":
             label = memory.get("label")
@@ -188,7 +209,13 @@ def _parse_memory_insert(memory: dict[str, Any], type: MemoryType) -> LabeledMem
                 raise ValueError(
                     f"The following column names are reserved: {', '.join(FORBIDDEN_METADATA_COLUMN_NAMES)}"
                 )
-            return {"value": value, "label": label, "source_id": source_id, "metadata": metadata}
+            return {
+                "value": value,
+                "label": label,
+                "source_id": source_id,
+                "partition_id": partition_id,
+                "metadata": metadata,
+            }
         case "SCORED":
             score = memory.get("score")
             if score is not None and not isinstance(score, (int, float)):
@@ -198,7 +225,13 @@ def _parse_memory_insert(memory: dict[str, Any], type: MemoryType) -> LabeledMem
                 raise ValueError(
                     f"The following column names are reserved: {', '.join(FORBIDDEN_METADATA_COLUMN_NAMES)}"
                 )
-            return {"value": value, "score": score, "source_id": source_id, "metadata": metadata}
+            return {
+                "value": value,
+                "score": score,
+                "source_id": source_id,
+                "partition_id": partition_id,
+                "metadata": metadata,
+            }
 
 
 def _parse_memory_update(update: dict[str, Any], type: MemoryType) -> LabeledMemoryUpdate | ScoredMemoryUpdate:
@@ -213,9 +246,15 @@ def _parse_memory_update(update: dict[str, Any], type: MemoryType) -> LabeledMem
             raise ValueError("value must be a string or unset")
         payload["value"] = update["value"]
     if "source_id" in update:
-        if not isinstance(update["source_id"], str):
-            raise ValueError("source_id must be a string or unset")
-        payload["source_id"] = update["source_id"]
+        source_id = update["source_id"]
+        if source_id is not None and not isinstance(source_id, str):
+            raise ValueError("source_id must be a string or None")
+        payload["source_id"] = source_id
+    if "partition_id" in update:
+        partition_id = update["partition_id"]
+        if partition_id is not None and not isinstance(partition_id, str):
+            raise ValueError("partition_id must be a string or None")
+        payload["partition_id"] = partition_id
     match type:
         case "LABELED":
             payload = cast(LabeledMemoryUpdate, payload)
@@ -249,6 +288,7 @@ class MemoryBase(ABC):
     value: str
     embedding: list[float]
     source_id: str | None
+    partition_id: str | None
     created_at: datetime
     updated_at: datetime
     metadata: dict[str, str | float | int | bool | None]
@@ -280,6 +320,7 @@ class MemoryBase(ABC):
         self.value = cast(str, memory["value"])
         self.embedding = memory["embedding"]
         self.source_id = memory["source_id"]
+        self.partition_id = memory["partition_id"]
         self.created_at = datetime.fromisoformat(memory["created_at"])
         self.updated_at = datetime.fromisoformat(memory["updated_at"])
         self.metadata = memory["metadata"]
@@ -292,11 +333,116 @@ class MemoryBase(ABC):
             raise AttributeError(f"{key} is not a valid attribute")
         return self.metadata[key]
 
+    def _convert_to_classification_prediction(
+        self,
+        prediction: LabelPredictionWithMemoriesAndFeedback,
+        *,
+        memoryset: LabeledMemoryset,
+        model: ClassificationModel,
+    ) -> ClassificationPrediction:
+        """
+        Convert internal prediction TypedDict to ClassificationPrediction object.
+        """
+        input_value = prediction.get("input_value")
+        input_value_str: str | None = None
+        if input_value is not None:
+            input_value_str = input_value.decode("utf-8") if isinstance(input_value, bytes) else input_value
+
+        return ClassificationPrediction(
+            prediction_id=prediction["prediction_id"],
+            label=prediction.get("label"),
+            label_name=prediction.get("label_name"),
+            score=None,
+            confidence=prediction["confidence"],
+            anomaly_score=prediction["anomaly_score"],
+            memoryset=memoryset,
+            model=model,
+            telemetry=prediction,
+            logits=prediction.get("logits"),
+            input_value=input_value_str,
+        )
+
+    def _convert_to_regression_prediction(
+        self,
+        prediction: ScorePredictionWithMemoriesAndFeedback,
+        *,
+        memoryset: ScoredMemoryset,
+        model: RegressionModel,
+    ) -> RegressionPrediction:
+        """
+        Convert internal prediction TypedDict to RegressionPrediction object.
+        """
+        input_value = prediction.get("input_value")
+        input_value_str: str | None = None
+        if input_value is not None:
+            input_value_str = input_value.decode("utf-8") if isinstance(input_value, bytes) else input_value
+
+        return RegressionPrediction(
+            prediction_id=prediction["prediction_id"],
+            label=None,
+            label_name=None,
+            score=prediction.get("score"),
+            confidence=prediction["confidence"],
+            anomaly_score=prediction["anomaly_score"],
+            memoryset=memoryset,
+            model=model,
+            telemetry=prediction,
+            logits=None,
+            input_value=input_value_str,
+        )
+
+    def feedback(self) -> dict[str, list[bool] | list[float]]:
+        """
+        Get feedback metrics computed from predictions that used this memory.
+
+        Returns a dictionary where:
+        - Keys are feedback category names
+        - Values are lists of feedback values (you may want to look at mean on the raw data)
+        """
+        # Collect all feedbacks by category, paginating through all predictions
+        feedback_by_category: dict[str, list[bool] | list[float]] = {}
+        batch_size = 500
+        offset = 0
+
+        while True:
+            predictions_batch = self.predictions(limit=batch_size, offset=offset)
+
+            if not predictions_batch:
+                break
+
+            for prediction in predictions_batch:
+                telemetry = prediction._telemetry
+                if "feedbacks" not in telemetry:
+                    continue
+
+                for fb in telemetry["feedbacks"]:
+                    category_name = fb["category_name"]
+                    value = fb["value"]
+                    # Convert BINARY (1/0) to boolean, CONTINUOUS to float
+                    if fb["category_type"] == "BINARY":
+                        value = bool(value)
+                        if category_name not in feedback_by_category:
+                            feedback_by_category[category_name] = []
+                        cast(list[bool], feedback_by_category[category_name]).append(value)
+                    else:
+                        value = float(value)
+                        if category_name not in feedback_by_category:
+                            feedback_by_category[category_name] = []
+                        cast(list[float], feedback_by_category[category_name]).append(value)
+
+            if len(predictions_batch) < batch_size:
+                break
+
+            offset += batch_size
+
+        return feedback_by_category
+
     def _update(
         self,
         *,
         value: str = UNSET,
         source_id: str | None = UNSET,
+        partition_id: str | None = UNSET,
         **metadata: None | bool | float | int | str,
     ) -> Self:
         client = OrcaClient._resolve_client()
@@ -307,6 +453,7 @@ class MemoryBase(ABC):
                 {"memory_id": self.memory_id}
                 | ({"value": value} if value is not UNSET else {})
                 | ({"source_id": source_id} if source_id is not UNSET else {})
+                | ({"partition_id": partition_id} if partition_id is not UNSET else {})
                 | {k: v for k, v in metadata.items() if v is not UNSET},
                 type=self.memory_type,
             ),
@@ -322,6 +469,7 @@ class MemoryBase(ABC):
             "value": self.value,
             "embedding": self.embedding,
             "source_id": self.source_id,
+            "partition_id": self.partition_id,
             "created_at": self.created_at,
             "updated_at": self.updated_at,
             "metadata": self.metadata,
@@ -346,6 +494,7 @@ class LabeledMemory(MemoryBase):
         label_name: Human-readable name of the label, automatically populated from the
             [`LabeledMemoryset.label_names`][orca_sdk.LabeledMemoryset]
         source_id: Optional unique identifier of the memory in a system of reference
+        partition_id: Optional identifier of the partition the memory belongs to
         metrics: Metrics about the memory, generated when running an analysis on the
             [`LabeledMemoryset`][orca_sdk.LabeledMemoryset]
         metadata: Metadata associated with the memory that is not used in the model. Metadata
@@ -384,6 +533,7 @@ class LabeledMemory(MemoryBase):
             + f"label: {('<' + self.label_name + ': ' + str(self.label) + '>') if self.label_name else str(self.label)}"
             + f", value: '{self.value[:100] + '...' if isinstance(self.value, str) and len(self.value) > 100 else self.value}'"
             + (f", source_id: '{self.source_id}'" if self.source_id is not None else "")
+            + (f", partition_id: '{self.partition_id}'" if self.partition_id is not None else "")
             + " })"
         )
 
@@ -396,6 +546,7 @@ class LabeledMemory(MemoryBase):
         value: str = UNSET,
         label: int | None = UNSET,
         source_id: str | None = UNSET,
+        partition_id: str | None = UNSET,
         **metadata: None | bool | float | int | str,
     ) -> LabeledMemory:
         """
@@ -408,14 +559,84 @@ class LabeledMemory(MemoryBase):
             value: New value of the memory
             label: New label of the memory
             source_id: New source ID of the memory
+            partition_id: New partition ID of the memory
             **metadata: New values for metadata properties
 
         Returns:
             The updated memory
         """
-        self._update(value=value, label=label, source_id=source_id, **metadata)
+        self._update(value=value, label=label, source_id=source_id, partition_id=partition_id, **metadata)
         return self
 
+    def predictions(
+        self,
+        limit: int = 100,
+        offset: int = 0,
+        tag: str | None = None,
+        sort: list[tuple[Literal["anomaly_score", "confidence", "timestamp"], Literal["asc", "desc"]]] = [],
+        expected_label_match: bool | None = None,
+    ) -> list[ClassificationPrediction]:
+        """
+        Get classification predictions that used this memory.
+
+        Args:
+            limit: Maximum number of predictions to return (default: 100)
+            offset: Number of predictions to skip for pagination (default: 0)
+            tag: Optional tag filter to only include predictions with this tag
+            sort: List of (field, direction) tuples for sorting results.
+                Valid fields: "anomaly_score", "confidence", "timestamp".
+                Valid directions: "asc", "desc"
+            expected_label_match: Filter by prediction correctness:
+                - True: only return correct predictions (label == expected_label)
+                - False: only return incorrect predictions (label != expected_label)
+                - None: return all predictions (default)
+
+        Returns:
+            List of ClassificationPrediction objects that used this memory
+        """
+
+        client = OrcaClient._resolve_client()
+        predictions_data = client.POST(
+            "/telemetry/prediction",
+            json={
+                "memory_id": self.memory_id,
+                "limit": limit,
+                "offset": offset,
+                "sort": [list(sort_item) for sort_item in sort],
+                "tag": tag,
+                "expected_label_match": expected_label_match,
+            },
+        )
+
+        # Filter to only classification predictions and convert to ClassificationPrediction objects
+        classification_predictions = [
+            cast(LabelPredictionWithMemoriesAndFeedback, p) for p in predictions_data if "label" in p
+        ]
+
+        from .classification_model import ClassificationModel
+
+        memorysets: dict[str, LabeledMemoryset] = {}
+        models: dict[str, ClassificationModel] = {}
+
+        def resolve_memoryset(memoryset_id: str) -> LabeledMemoryset:
+            if memoryset_id not in memorysets:
+                memorysets[memoryset_id] = LabeledMemoryset.open(memoryset_id)
+            return memorysets[memoryset_id]
+
+        def resolve_model(model_id: str) -> ClassificationModel:
+            if model_id not in models:
+                models[model_id] = ClassificationModel.open(model_id)
+            return models[model_id]
+
+        return [
+            self._convert_to_classification_prediction(
+                p,
+                memoryset=resolve_memoryset(p["memoryset_id"]),
+                model=resolve_model(p["model_id"]),
+            )
+            for p in classification_predictions
+        ]
+
     def to_dict(self) -> dict[str, Any]:
         """
         Convert the memory to a dictionary
@@ -441,6 +662,7 @@ class LabeledMemoryLookup(LabeledMemory):
         label_name: Human-readable name of the label, automatically populated from the
             [`LabeledMemoryset.label_names`][orca_sdk.LabeledMemoryset]
         source_id: Optional unique identifier of the memory in a system of reference
+        partition_id: Optional identifier of the partition the memory belongs to
         metrics: Metrics about the memory, generated when running an analysis on the
             [`LabeledMemoryset`][orca_sdk.LabeledMemoryset]
         metadata: Metadata associated with the memory that is not used in the model. Metadata
@@ -457,7 +679,11 @@ class LabeledMemoryLookup(LabeledMemory):
     lookup_score: float
     attention_weight: float | None
 
-    def __init__(self, memoryset_id: str, memory_lookup: LabeledMemoryLookupResponse | LabelPredictionMemoryLookup):
+    def __init__(
+        self,
+        memoryset_id: str,
+        memory_lookup: LabeledMemoryLookupResponse | LabelPredictionMemoryLookup,
+    ):
         # for internal use only, do not document
         super().__init__(memoryset_id, memory_lookup)
         self.lookup_score = memory_lookup["lookup_score"]
@@ -471,6 +697,7 @@ class LabeledMemoryLookup(LabeledMemory):
             + (f", attention_weight: {self.attention_weight:.2f}" if self.attention_weight is not None else "")
             + f", value: '{self.value[:100] + '...' if isinstance(self.value, str) and len(self.value) > 100 else self.value}'"
             + (f", source_id: '{self.source_id}'" if self.source_id is not None else "")
+            + (f", partition_id: '{self.partition_id}'" if self.partition_id is not None else "")
             + " })"
         )
 
@@ -485,6 +712,7 @@ class ScoredMemory(MemoryBase):
             with the [`ScoredMemoryset.embedding_model`][orca_sdk.ScoredMemoryset]
         score: Score of the memory
         source_id: Optional unique identifier of the memory in a system of reference
+        partition_id: Optional identifier of the partition the memory belongs to
         metrics: Metrics about the memory, generated when running an analysis on the
             [`ScoredMemoryset`][orca_sdk.ScoredMemoryset]
         metadata: Metadata associated with the memory that is not used in the model. Metadata
@@ -521,6 +749,7 @@ class ScoredMemory(MemoryBase):
             + f"score: {self.score:.2f}"
             + f", value: '{self.value[:100] + '...' if isinstance(self.value, str) and len(self.value) > 100 else self.value}'"
             + (f", source_id: '{self.source_id}'" if self.source_id is not None else "")
+            + (f", partition_id: '{self.partition_id}'" if self.partition_id is not None else "")
             + " })"
         )
 
@@ -533,6 +762,7 @@ class ScoredMemory(MemoryBase):
         value: str = UNSET,
         score: float | None = UNSET,
         source_id: str | None = UNSET,
+        partition_id: str | None = UNSET,
         **metadata: None | bool | float | int | str,
     ) -> ScoredMemory:
         """
@@ -550,9 +780,78 @@ class ScoredMemory(MemoryBase):
         Returns:
             The updated memory
         """
-        self._update(value=value, score=score, source_id=source_id, **metadata)
+        self._update(value=value, score=score, source_id=source_id, partition_id=partition_id, **metadata)
         return self
 
+    def predictions(
+        self,
+        limit: int = 100,
+        offset: int = 0,
+        tag: str | None = None,
+        sort: list[tuple[Literal["anomaly_score", "confidence", "timestamp"], Literal["asc", "desc"]]] = [],
+        expected_label_match: bool | None = None,
+    ) -> list[RegressionPrediction]:
+        """
+        Get regression predictions that used this memory.
+
+        Args:
+            limit: Maximum number of predictions to return (default: 100)
+            offset: Number of predictions to skip for pagination (default: 0)
+            tag: Optional tag filter to only include predictions with this tag
+            sort: List of (field, direction) tuples for sorting results.
+                Valid fields: "anomaly_score", "confidence", "timestamp".
+                Valid directions: "asc", "desc"
+            expected_label_match: Filter by prediction correctness:
+                - True: only return correct predictions (score close to expected_score)
+                - False: only return incorrect predictions (score differs from expected_score)
+                - None: return all predictions (default)
+                Note: For regression, "correctness" is based on score proximity to expected_score.
+
+        Returns:
+            List of RegressionPrediction objects that used this memory
+        """
+        client = OrcaClient._resolve_client()
+        predictions_data = client.POST(
+            "/telemetry/prediction",
+            json={
+                "memory_id": self.memory_id,
+                "limit": limit,
+                "offset": offset,
+                "sort": [list(sort_item) for sort_item in sort],
+                "tag": tag,
+                "expected_label_match": expected_label_match,
+            },
+        )
+
+        # Filter to only regression predictions and convert to RegressionPrediction objects
+        regression_predictions = [
+            cast(ScorePredictionWithMemoriesAndFeedback, p) for p in predictions_data if "score" in p
+        ]
+
+        from .regression_model import RegressionModel
+
+        memorysets: dict[str, ScoredMemoryset] = {}
+        models: dict[str, RegressionModel] = {}
+
+        def resolve_memoryset(memoryset_id: str) -> ScoredMemoryset:
+            if memoryset_id not in memorysets:
+                memorysets[memoryset_id] = ScoredMemoryset.open(memoryset_id)
+            return memorysets[memoryset_id]
+
+        def resolve_model(model_id: str) -> RegressionModel:
+            if model_id not in models:
+                models[model_id] = RegressionModel.open(model_id)
+            return models[model_id]
+
+        return [
+            self._convert_to_regression_prediction(
+                p,
+                memoryset=resolve_memoryset(p["memoryset_id"]),
+                model=resolve_model(p["model_id"]),
+            )
+            for p in regression_predictions
+        ]
+
     def to_dict(self) -> dict[str, Any]:
         """
         Convert the memory to a dictionary
@@ -575,6 +874,7 @@ class ScoredMemoryLookup(ScoredMemory):
             with the [`ScoredMemoryset.embedding_model`][orca_sdk.ScoredMemoryset]
         score: Score of the memory
         source_id: Optional unique identifier of the memory in a system of reference
+        partition_id: Optional identifier of the partition the memory belongs to
         metrics: Metrics about the memory, generated when running an analysis on the
             [`ScoredMemoryset`][orca_sdk.ScoredMemoryset]
         memory_id: The unique identifier for the memory, automatically generated on insert
@@ -589,7 +889,11 @@ class ScoredMemoryLookup(ScoredMemory):
     lookup_score: float
     attention_weight: float | None
 
-    def __init__(self, memoryset_id: str, memory_lookup: ScoredMemoryLookupResponse | ScorePredictionMemoryLookup):
+    def __init__(
+        self,
+        memoryset_id: str,
+        memory_lookup: ScoredMemoryLookupResponse | ScorePredictionMemoryLookup,
+    ):
         # for internal use only, do not document
         super().__init__(memoryset_id, memory_lookup)
         self.lookup_score = memory_lookup["lookup_score"]
@@ -602,6 +906,7 @@ class ScoredMemoryLookup(ScoredMemory):
             + f", lookup_score: {self.lookup_score:.2f}"
             + f", value: '{self.value[:100] + '...' if isinstance(self.value, str) and len(self.value) > 100 else self.value}'"
             + (f", source_id: '{self.source_id}'" if self.source_id is not None else "")
+            + (f", partition_id: '{self.partition_id}'" if self.partition_id is not None else "")
             + " })"
         )
 
@@ -727,6 +1032,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         label_column: str | None = None,
         score_column: str | None = None,
         source_id_column: str | None = None,
+        partition_id_column: str | None = None,
         description: str | None = None,
         label_names: list[str] | None = None,
         max_seq_length_override: int | None = None,
@@ -737,12 +1043,14 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         if_exists: CreateMode = "error",
         background: bool = False,
         hidden: bool = False,
+        subsample: int | float | None = None,
+        memory_type: MemoryType | None = None,
     ) -> Self | Job[Self]:
         """
         Create a new memoryset in the OrcaCloud
 
         All columns from the datasource that are not specified in the `value_column`,
-        `label_column`, or `source_id_column` will be stored as metadata in the memoryset.
+        `label_column`, `source_id_column`, or `partition_id_column` will be stored as metadata in the memoryset.
 
         Params:
             name: Name for the new memoryset (must be unique)
@@ -750,18 +1058,20 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             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.
             value_column: Name of the column in the datasource that contains the memory values
-            label_column: Name of the column in the datasource that contains the memory labels,
-                these must be contiguous integers starting from 0
+            label_column: Name of the column in the datasource that contains the memory labels.
+                Must contain categorical values as integers or strings. String labels will be
+                converted to integers with the unique strings extracted as `label_names`
             score_column: Name of the column in the datasource that contains the memory scores
             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
             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, must match
-                the number of labels in the `label_column`. Will be automatically inferred if a
-                [Dataset][datasets.Dataset] with a [`ClassLabel`][datasets.ClassLabel] feature for
-                labels is used as the datasource
+                the number of labels in the `label_column`. Will be automatically inferred if string
+                labels are provided or if a [Dataset][datasets.Dataset] with a
+                [`ClassLabel`][datasets.ClassLabel] feature for labels is used as the datasource
             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
@@ -775,7 +1085,10 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 `"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
-
+            subsample: Optional number (int) of rows to insert or fraction (float in (0, 1]) of the
+                datasource to insert. Use to limit the size of the initial memoryset.
+            memory_type: Type of memoryset to create, defaults to `"LABELED"` if `label_column` is provided,
+                and `"SCORED"` if `score_column` is provided, must be specified for other cases.
         Returns:
             Handle to the new memoryset in the OrcaCloud
 
@@ -786,9 +1099,6 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         if embedding_model is None:
             embedding_model = PretrainedEmbeddingModel.GTE_BASE
 
-        if label_column is None and score_column is None:
-            raise ValueError("label_column or score_column must be provided")
-
         existing = cls._handle_if_exists(
             name,
             if_exists=if_exists,
@@ -806,6 +1116,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             "datasource_score_column": score_column,
             "datasource_value_column": value_column,
             "datasource_source_id_column": source_id_column,
+            "datasource_partition_id_column": partition_id_column,
             "label_names": label_names,
             "max_seq_length_override": max_seq_length_override,
             "remove_duplicates": remove_duplicates,
@@ -813,6 +1124,10 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             "index_params": index_params,
             "hidden": hidden,
         }
+        if memory_type is not None:
+            payload["memory_type"] = memory_type
+        if subsample is not None:
+            payload["subsample"] = subsample
         if prompt is not None:
             payload["prompt"] = prompt
         if isinstance(embedding_model, PretrainedEmbeddingModel):
@@ -823,7 +1138,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             raise ValueError("Invalid embedding model")
         client = OrcaClient._resolve_client()
         response = client.POST("/memoryset", json=payload)
-        job = Job(response["insertion_task_id"], lambda: cls.open(response["id"]))
+        job = Job(response["insertion_job_id"], lambda: cls.open(response["id"]))
         return job if background else job.result()
 
     @overload
@@ -918,7 +1233,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         appended with `_datasource` and use that as the datasource for the memoryset.
 
         All properties that are not specified to be used as `value_column`, `label_column`, or
-        `source_id_column` will be stored as metadata in the memoryset.
+        `source_id_column`, or `partition_id_column` will be stored as metadata in the memoryset.
 
         Params:
             name: Name for the new memoryset (must be unique)
@@ -988,7 +1303,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         appended with `_datasource` and use that as the datasource for the memoryset.
 
         All properties that are not specified to be used as `value_column`, `label_column`, or
-        `source_id_column` will be stored as metadata in the memoryset.
+        `source_id_column`, or `partition_id_column` will be stored as metadata in the memoryset.
 
         Params:
             name: Name for the new memoryset (must be unique)
@@ -1060,7 +1375,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         appended with `_datasource` and use that as the datasource for the memoryset.
 
         All columns from the datasource that are not specified in the `value_column`,
-        `label_column`, or `source_id_column` will be stored as metadata in the memoryset.
+        `label_column`, `source_id_column`, or `partition_id_column` will be stored as metadata in the memoryset.
 
         Params:
             name: Name for the new memoryset (must be unique)
@@ -1133,7 +1448,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         appended with `_datasource` and use that as the datasource for the memoryset.
 
         All columns that are not specified to be used as `value_column`, `label_column`, or
-        `source_id_column` will be stored as metadata in the memoryset.
+        `source_id_column`, or `partition_id_column` will be stored as metadata in the memoryset.
 
         Params:
             name: Name for the new memoryset (must be unique)
@@ -1199,7 +1514,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         appended with `_datasource` and use that as the datasource for the memoryset.
 
         All columns that are not specified to be used as `value_column`, `label_column`, or
-        `source_id_column` will be stored as metadata in the memoryset.
+        `source_id_column`, or `partition_id_column` will be stored as metadata in the memoryset.
 
         Params:
             name: Name for the new memoryset (must be unique)
@@ -1267,7 +1582,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         appended with `_datasource` and use that as the datasource for the memoryset.
 
         All columns from the datasource that are not specified in the `value_column`,
-        `label_column`, or `source_id_column` will be stored as metadata in the memoryset.
+        `label_column`, `source_id_column`, or `partition_id_column` will be stored as metadata in the memoryset.
 
         Params:
             name: Name for the new memoryset (must be unique)
@@ -1516,7 +1831,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         client = OrcaClient._resolve_client()
         metadata = client.POST("/memoryset/{name_or_id}/clone", params={"name_or_id": self.id}, json=payload)
         job = Job(
-            metadata["insertion_task_id"],
+            metadata["insertion_job_id"],
             lambda: self.open(metadata["id"]),
         )
         return job if background else job.result()
@@ -1595,15 +1910,43 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             raise ValueError(f"Invalid index type: {type(index)}")
 
     @overload
-    def search(self, query: str, *, count: int = 1, prompt: str | None = None) -> list[MemoryLookupT]:
+    def search(
+        self,
+        query: str,
+        *,
+        count: int = 1,
+        prompt: str | None = None,
+        partition_id: str | None = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
+    ) -> list[MemoryLookupT]:
         pass
 
     @overload
-    def search(self, query: list[str], *, count: int = 1, prompt: str | None = None) -> list[list[MemoryLookupT]]:
+    def search(
+        self,
+        query: list[str],
+        *,
+        count: int = 1,
+        prompt: str | None = None,
+        partition_id: str | None = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
+    ) -> list[list[MemoryLookupT]]:
         pass
 
     def search(
-        self, query: str | list[str], *, count: int = 1, prompt: str | None = None
+        self,
+        query: str | list[str],
+        *,
+        count: int = 1,
+        prompt: str | None = None,
+        partition_id: str | None = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> list[MemoryLookupT] | list[list[MemoryLookupT]]:
         """
         Search for memories that are semantically similar to the query
@@ -1613,7 +1956,12 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             count: Number of memories to return for each query
             prompt: Optional prompt for query embedding during search.
                 If not provided, the memoryset's default query prompt will be used if available.
-
+            partition_id: Optional partition ID to filter memories by
+            partition_filter_mode: How to filter partitions when searching for memories
+                - "ignore_partitions": Ignore partitions
+                - "include_global": Include global memories
+                - "exclude_global": Exclude global memories
+                - "only_global": Only include global memories
         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
@@ -1653,6 +2001,8 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 "query": query if isinstance(query, list) else [query],
                 "count": count,
                 "prompt": prompt,
+                "partition_id": partition_id,
+                "partition_filter_mode": partition_filter_mode,
             },
         )
         lookups = [
@@ -1678,6 +2028,10 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         filters: list[FilterItemTuple] = [],
         with_feedback_metrics: bool = False,
         sort: list[TelemetrySortItem] | None = None,
+        partition_id: str | None = None,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> list[MemoryT]:
         """
         Query the memoryset for memories that match the filters
@@ -1703,6 +2057,13 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         ]
 
         if with_feedback_metrics:
+            if partition_id:
+                raise ValueError("Partition ID is not supported when with_feedback_metrics is True")
+            if partition_filter_mode != "include_global":
+                raise ValueError(
+                    f"Partition filter mode {partition_filter_mode} is not supported when with_feedback_metrics is True. Only 'include_global' is supported."
+                )
+
             client = OrcaClient._resolve_client()
             response = client.POST(
                 "/telemetry/memories",
@@ -1736,6 +2097,8 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 "offset": offset,
                 "limit": limit,
                 "filters": cast(list[FilterItem], parsed_filters),
+                "partition_id": partition_id,
+                "partition_filter_mode": partition_filter_mode,
             },
         )
         return [
@@ -1786,8 +2149,8 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
 
         Examples:
             >>> memoryset.insert([
-            ...     {"value": "I am happy", "label": 1, "source_id": "user_123", "tag": "happy"},
-            ...     {"value": "I am sad", "label": 0, "source_id": "user_124", "tag": "sad"},
+            ...     {"value": "I am happy", "label": 1, "source_id": "data_123", "partition_id": "user_1", "tag": "happy"},
+            ...     {"value": "I am sad", "label": 0, "source_id": "data_124", "partition_id": "user_1", "tag": "sad"},
             ... ])
         """
         client = OrcaClient._resolve_client()
@@ -1818,12 +2181,13 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 - `label`: Label of the memory
                 - `score`: Score of the memory
                 - `source_id`: Optional unique ID of the memory in a system of reference
+                - `partition_id`: Optional partition ID of the memory
                 - `...`: Any other metadata to store for the memory
 
         Examples:
             >>> await memoryset.ainsert([
-            ...     {"value": "I am happy", "label": 1, "source_id": "user_123", "tag": "happy"},
-            ...     {"value": "I am sad", "label": 0, "source_id": "user_124", "tag": "sad"},
+            ...     {"value": "I am happy", "label": 1, "source_id": "data_123", "partition_id": "user_1", "tag": "happy"},
+            ...     {"value": "I am sad", "label": 0, "source_id": "data_124", "partition_id": "user_1", "tag": "sad"},
             ... ])
         """
         client = OrcaAsyncClient._resolve_client()
@@ -1938,6 +2302,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 - `value`: Optional new value of the memory
                 - `label`: Optional new label of the memory
                 - `source_id`: Optional new source ID of the memory
+                - `partition_id`: Optional new partition ID of the memory
                 - `...`: Optional new values for metadata properties
 
         Returns:
@@ -2075,6 +2440,9 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         lookup_count: int = 15,
         clear_metrics: bool = False,
         background: Literal[True],
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> Job[MemorysetMetrics]:
         pass
 
@@ -2085,6 +2453,9 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         lookup_count: int = 15,
         clear_metrics: bool = False,
         background: Literal[False] = False,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> MemorysetMetrics:
         pass
 
@@ -2094,6 +2465,9 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         lookup_count: int = 15,
         clear_metrics: bool = False,
         background: bool = False,
+        partition_filter_mode: Literal[
+            "ignore_partitions", "include_global", "exclude_global", "only_global"
+        ] = "include_global",
     ) -> Job[MemorysetMetrics] | MemorysetMetrics:
         """
         Run analyses on the memoryset to find duplicates, clusters, mislabelings, and more
@@ -2114,6 +2488,11 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
 
             lookup_count: Number of memories to lookup for each memory in the memoryset
             clear_metrics: Whether to clear any existing metrics from the memories before running the analysis
+            partition_filter_mode: How to filter partitions when running the analysis
+                - "ignore_partitions": Ignore partitions
+                - "include_global": Include global memories
+                - "exclude_global": Exclude global memories
+                - "only_global": Only include global memories
 
         Returns:
             dictionary with aggregate metrics for each analysis that was run
@@ -2183,17 +2562,18 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 "configs": configs,
                 "lookup_count": lookup_count,
                 "clear_metrics": clear_metrics,
+                "partition_filter_mode": partition_filter_mode,
             },
         )
 
         def get_analysis_result():
             client = OrcaClient._resolve_client()
             return client.GET(
-                "/memoryset/{name_or_id}/analysis/{analysis_task_id}",
-                params={"name_or_id": self.id, "analysis_task_id": analysis["task_id"]},
+                "/memoryset/{name_or_id}/analysis/{analysis_job_id}",
+                params={"name_or_id": self.id, "analysis_job_id": analysis["job_id"]},
             )["results"]
 
-        job = Job(analysis["task_id"], get_analysis_result)
+        job = Job(analysis["job_id"], get_analysis_result)
         return job if background else job.result()
 
     def get_potential_duplicate_groups(self) -> list[list[MemoryT]]:
@@ -2241,8 +2621,9 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         *,
         embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
         value_column: str = "value",
-        label_column: str = "label",
+        label_column: str | None = "label",
         source_id_column: str | None = None,
+        partition_id_column: str | None = None,
         description: str | None = None,
         label_names: list[str] | None = None,
         max_seq_length_override: int | None = None,
@@ -2253,6 +2634,7 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         if_exists: CreateMode = "error",
         background: Literal[True],
         hidden: bool = False,
+        subsample: int | float | None = None,
     ) -> Job[Self]:
         pass
 
@@ -2265,8 +2647,9 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         *,
         embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
         value_column: str = "value",
-        label_column: str = "label",
+        label_column: str | None = "label",
         source_id_column: str | None = None,
+        partition_id_column: str | None = None,
         description: str | None = None,
         label_names: list[str] | None = None,
         max_seq_length_override: int | None = None,
@@ -2277,6 +2660,7 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         if_exists: CreateMode = "error",
         background: Literal[False] = False,
         hidden: bool = False,
+        subsample: int | float | None = None,
     ) -> Self:
         pass
 
@@ -2288,8 +2672,9 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         *,
         embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
         value_column: str = "value",
-        label_column: str = "label",
+        label_column: str | None = "label",
         source_id_column: str | None = None,
+        partition_id_column: str | None = None,
         description: str | None = None,
         label_names: list[str] | None = None,
         max_seq_length_override: int | None = None,
@@ -2300,12 +2685,13 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         if_exists: CreateMode = "error",
         background: bool = False,
         hidden: bool = False,
+        subsample: int | float | None = None,
     ) -> Self | Job[Self]:
         """
         Create a new labeled memoryset in the OrcaCloud
 
         All columns from the datasource that are not specified in the `value_column`,
-        `label_column`, or `source_id_column` will be stored as metadata in the memoryset.
+        `label_column`, `source_id_column`, or `partition_id_column` will be stored as metadata in the memoryset.
 
         Params:
             name: Name for the new memoryset (must be unique)
@@ -2313,17 +2699,20 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
             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.
             value_column: Name of the column in the datasource that contains the memory values
-            label_column: Name of the column in the datasource that contains the memory labels,
-                these must be contiguous integers starting from 0
+            label_column: Name of the column in the datasource that contains the memory labels.
+                Must contain categorical values as integers or strings. String labels will be
+                converted to integers with the unique strings extracted as `label_names`. To create
+                a memoryset with all none labels, set to `None`.
             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
             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, must match
-                the number of labels in the `label_column`. Will be automatically inferred if a
-                [Dataset][datasets.Dataset] with a [`ClassLabel`][datasets.ClassLabel] feature for
-                labels is used as the datasource
+                the number of labels in the `label_column`. Will be automatically inferred if string
+                labels are provided or if a [Dataset][datasets.Dataset] with a
+                [`ClassLabel`][datasets.ClassLabel] feature for labels is used as the datasource
             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
@@ -2353,6 +2742,7 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
             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,
@@ -2363,6 +2753,8 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
             if_exists=if_exists,
             background=background,
             hidden=hidden,
+            subsample=subsample,
+            memory_type="LABELED",
         )
 
     def display_label_analysis(self):
@@ -2405,8 +2797,9 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         *,
         embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
         value_column: str = "value",
-        score_column: str = "score",
+        score_column: str | None = "score",
         source_id_column: str | None = None,
+        partition_id_column: str | None = None,
         description: str | None = None,
         max_seq_length_override: int | None = None,
         prompt: str | None = None,
@@ -2416,6 +2809,7 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         if_exists: CreateMode = "error",
         background: Literal[True],
         hidden: bool = False,
+        subsample: int | float | None = None,
     ) -> Job[Self]:
         pass
 
@@ -2427,9 +2821,10 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         datasource: Datasource,
         *,
         embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
-        score_column: str = "score",
+        score_column: str | None = "score",
         value_column: str = "value",
         source_id_column: str | None = None,
+        partition_id_column: str | None = None,
         description: str | None = None,
         max_seq_length_override: int | None = None,
         prompt: str | None = None,
@@ -2439,6 +2834,7 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         if_exists: CreateMode = "error",
         background: Literal[False] = False,
         hidden: bool = False,
+        subsample: int | float | None = None,
     ) -> Self:
         pass
 
@@ -2450,8 +2846,9 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         *,
         embedding_model: FinetunedEmbeddingModel | PretrainedEmbeddingModel | None = None,
         value_column: str = "value",
-        score_column: str = "score",
+        score_column: str | None = "score",
         source_id_column: str | None = None,
+        partition_id_column: str | None = None,
         description: str | None = None,
         max_seq_length_override: int | None = None,
         prompt: str | None = None,
@@ -2461,12 +2858,13 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         if_exists: CreateMode = "error",
         background: bool = False,
         hidden: bool = False,
+        subsample: int | float | None = None,
     ) -> Self | Job[Self]:
         """
         Create a new scored memoryset in the OrcaCloud
 
         All columns from the datasource that are not specified in the `value_column`,
-        `score_column`, or `source_id_column` will be stored as metadata in the memoryset.
+        `score_column`, `source_id_column`, or `partition_id_column` will be stored as metadata in the memoryset.
 
         Params:
             name: Name for the new memoryset (must be unique)
@@ -2474,9 +2872,11 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
             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.
             value_column: Name of the column in the datasource that contains the memory values
-            score_column: Name of the column in the datasource that contains the memory scores
+            score_column: Name of the column in the datasource that contains the memory scores. Must
+                contain numerical values. To create a memoryset with all none scores, set to `None`.
             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
             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.
@@ -2508,6 +2908,7 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
             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,
@@ -2517,4 +2918,6 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
             if_exists=if_exists,
             background=background,
             hidden=hidden,
+            subsample=subsample,
+            memory_type="SCORED",
         )