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

orca_sdk/memoryset.py

@@ -114,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",
@@ -194,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")
@@ -206,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)):
@@ -216,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:
@@ -231,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)
@@ -267,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]
@@ -298,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"]
@@ -419,6 +442,7 @@ class MemoryBase(ABC):
         *,
         value: str = UNSET,
         source_id: str | None = UNSET,
+        partition_id: str | None = UNSET,
         **metadata: None | bool | float | int | str,
     ) -> Self:
         client = OrcaClient._resolve_client()
@@ -429,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,
             ),
@@ -444,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,
@@ -468,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
@@ -506,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 "")
             + " })"
         )
 
@@ -518,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:
         """
@@ -530,12 +559,13 @@ 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(
@@ -632,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
@@ -666,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 "")
             + " })"
         )
 
@@ -680,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
@@ -716,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 "")
             + " })"
         )
 
@@ -728,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:
         """
@@ -745,7 +780,7 @@ 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(
@@ -839,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
@@ -870,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 "")
             + " })"
         )
 
@@ -995,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,
@@ -1012,7 +1050,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
         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)
@@ -1026,6 +1064,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             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.
@@ -1077,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,
@@ -1193,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)
@@ -1263,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)
@@ -1335,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)
@@ -1408,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)
@@ -1474,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)
@@ -1542,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)
@@ -1870,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
@@ -1888,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
@@ -1928,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 = [
@@ -1953,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
@@ -1978,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",
@@ -2011,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 [
@@ -2061,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()
@@ -2093,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()
@@ -2213,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:
@@ -2350,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
 
@@ -2360,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
 
@@ -2369,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
@@ -2389,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
@@ -2458,6 +2562,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 "configs": configs,
                 "lookup_count": lookup_count,
                 "clear_metrics": clear_metrics,
+                "partition_filter_mode": partition_filter_mode,
             },
         )
 
@@ -2518,6 +2623,7 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         value_column: str = "value",
         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,
@@ -2543,6 +2649,7 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         value_column: str = "value",
         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,
@@ -2567,6 +2674,7 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         value_column: str = "value",
         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,
@@ -2583,7 +2691,7 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
         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)
@@ -2597,6 +2705,7 @@ class LabeledMemoryset(MemorysetBase[LabeledMemory, LabeledMemoryLookup]):
                 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.
@@ -2633,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,
@@ -2689,6 +2799,7 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         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,
         max_seq_length_override: int | None = None,
         prompt: str | None = None,
@@ -2713,6 +2824,7 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         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,
@@ -2736,6 +2848,7 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         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,
         max_seq_length_override: int | None = None,
         prompt: str | None = None,
@@ -2751,7 +2864,7 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
         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)
@@ -2763,6 +2876,7 @@ class ScoredMemoryset(MemorysetBase[ScoredMemory, ScoredMemoryLookup]):
                 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.
@@ -2794,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,