orca-sdk 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

orca_sdk/client.py

@@ -17,6 +17,7 @@ from typing import (
     Mapping,
     NotRequired,
     Self,
+    TypeAlias,
     TypedDict,
     cast,
     overload,
@@ -153,11 +154,14 @@ class ClusterMetrics(TypedDict):
     memory_count: int
 
 
-ColumnType = Literal["STRING", "FLOAT", "INT", "BOOL", "ENUM", "IMAGE", "OTHER"]
+ColumnType: TypeAlias = Literal["STRING", "FLOAT", "INT", "BOOL", "ENUM", "IMAGE", "OTHER"]
+"""
+The type of a column in a datasource
+"""
 
 
 class ConstraintViolationErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[409]
     constraint: str
 
 
@@ -168,6 +172,7 @@ class CountPredictionsRequest(TypedDict):
     start_timestamp: NotRequired[str | None]
     end_timestamp: NotRequired[str | None]
     memory_id: NotRequired[str | None]
+    expected_label_match: NotRequired[bool | None]
 
 
 class CreateApiKeyRequest(TypedDict):
@@ -230,7 +235,7 @@ class EmbeddingEvaluationRequest(TypedDict):
     weigh_memories: NotRequired[bool]
 
 
-EmbeddingFinetuningMethod = Literal["classification", "regression", "batch_triplet_loss"]
+EmbeddingFinetuningMethod: TypeAlias = Literal["classification", "regression", "batch_triplet_loss"]
 
 
 class FeedbackMetrics(TypedDict):
@@ -238,13 +243,55 @@ class FeedbackMetrics(TypedDict):
     count: int
 
 
-FeedbackType = Literal["CONTINUOUS", "BINARY"]
+FeedbackType: TypeAlias = Literal["CONTINUOUS", "BINARY"]
 
 
 class FilterItem(TypedDict):
-    field: list
-    op: Literal["==", "!=", ">", ">=", "<", "<=", "in", "not in", "like"]
+    field: (
+        tuple[
+            Literal[
+                "memory_id",
+                "value",
+                "label",
+                "metadata",
+                "source_id",
+                "partition_id",
+                "created_at",
+                "updated_at",
+                "edited_at",
+                "metrics",
+                "score",
+                "labels",
+            ]
+        ]
+        | tuple[Literal["metadata"], str]
+        | tuple[
+            Literal["metrics"],
+            Literal[
+                "cluster",
+                "embedding_2d",
+                "is_duplicate",
+                "duplicate_memory_ids",
+                "has_potential_duplicates",
+                "potential_duplicate_memory_ids",
+                "anomaly_score",
+                "neighbor_label_logits",
+                "neighbor_predicted_label",
+                "neighbor_predicted_label_ambiguity",
+                "neighbor_predicted_label_confidence",
+                "current_label_neighbor_confidence",
+                "normalized_neighbor_label_entropy",
+                "neighbor_predicted_label_matches_current_label",
+                "spread",
+                "uniformity",
+                "concept_id",
+                "subconcept_id",
+            ],
+        ]
+    )
+    op: Literal["==", "!=", ">", ">=", "<", "<=", "in", "not in", "like", "contains all", "contains any"]
     value: str | int | float | bool | list[str | None] | list[int] | list[float] | list[bool] | None
+    transform: NotRequired[Literal["length"]]
 
 
 class GetDatasourceRowCountRequest(TypedDict):
@@ -272,12 +319,17 @@ class HealthyResponse(TypedDict):
 
 
 class InternalServerErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[500]
     message: str
     request_id: str
 
 
-JobStatus = Literal["INITIALIZED", "DISPATCHED", "WAITING", "PROCESSING", "COMPLETED", "FAILED", "ABORTING", "ABORTED"]
+JobStatus: TypeAlias = Literal[
+    "INITIALIZED", "DISPATCHED", "WAITING", "PROCESSING", "COMPLETED", "FAILED", "ABORTING", "ABORTED"
+]
+"""
+Status of job in the job queue
+"""
 
 
 class JobStatusInfo(TypedDict):
@@ -344,7 +396,7 @@ class MemoryMetrics(TypedDict):
     has_potential_duplicates: NotRequired[bool]
     potential_duplicate_memory_ids: NotRequired[list[str] | None]
     cluster: NotRequired[int]
-    embedding_2d: NotRequired[list]
+    embedding_2d: NotRequired[tuple[float, float]]
     anomaly_score: NotRequired[float]
     neighbor_label_logits: NotRequired[list[float] | None]
     neighbor_predicted_label: NotRequired[int | None]
@@ -359,7 +411,7 @@ class MemoryMetrics(TypedDict):
     subconcept_id: NotRequired[int | None]
 
 
-MemoryType = Literal["LABELED", "SCORED"]
+MemoryType: TypeAlias = Literal["LABELED", "SCORED"]
 
 
 class MemorysetClassPatternsAnalysisConfig(TypedDict):
@@ -465,7 +517,7 @@ class MemorysetUpdate(TypedDict):
 
 
 class NotFoundErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[404]
     resource: (
         Literal[
             "org",
@@ -545,7 +597,7 @@ class PredictionFeedbackResult(TypedDict):
     new_category_ids: list[str]
 
 
-PredictionSort = list[list]
+PredictionSort: TypeAlias = list[tuple[Literal["timestamp", "confidence", "anomaly_score"], Literal["asc", "desc"]]]
 
 
 class PredictiveModelUpdate(TypedDict):
@@ -554,15 +606,18 @@ class PredictiveModelUpdate(TypedDict):
     locked: NotRequired[bool]
 
 
-PretrainedEmbeddingModelName = Literal[
+PretrainedEmbeddingModelName: TypeAlias = Literal[
     "CLIP_BASE", "GTE_BASE", "CDE_SMALL", "DISTILBERT", "GTE_SMALL", "MXBAI_LARGE", "E5_LARGE", "BGE_BASE", "GIST_LARGE"
 ]
+"""
+Names of pretrained embedding models that are supported by OrcaCloud
+"""
 
 
-RACHeadType = Literal["KNN", "MMOE", "FF", "BMMOE"]
+RACHeadType: TypeAlias = Literal["KNN", "MMOE", "FF", "BMMOE"]
 
 
-RARHeadType = Literal["MMOE", "KNN"]
+RARHeadType: TypeAlias = Literal["MMOE", "KNN"]
 
 
 class ROCCurve(TypedDict):
@@ -669,6 +724,7 @@ class ScorePredictionWithMemoriesAndFeedback(TypedDict):
     tags: list[str]
     explanation: str | None
     memory_id: str | None
+    is_in_dense_neighborhood: NotRequired[bool | None]
     feedbacks: list[PredictionFeedback]
 
 
@@ -740,7 +796,7 @@ class ScoredMemoryWithFeedbackMetrics(TypedDict):
 
 
 class ServiceUnavailableErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[503]
     service: str
 
 
@@ -752,7 +808,9 @@ class SubConceptMetrics(TypedDict):
     memory_count: int
 
 
-TelemetryField = list
+TelemetryField: TypeAlias = (
+    tuple[Literal["feedback_metrics"], str, Literal["avg", "count"]] | tuple[Literal["lookup"], Literal["count"]]
+)
 
 
 class TelemetryFilterItem(TypedDict):
@@ -767,11 +825,11 @@ class TelemetrySortOptions(TypedDict):
 
 
 class UnauthenticatedErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[401]
 
 
 class UnauthorizedErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[403]
     reason: str
 
 
@@ -792,7 +850,10 @@ class ValidationError(TypedDict):
     type: str
 
 
-WorkerStatus = Literal["IDLE", "BUSY", "DRAINING", "SHUTDOWN", "CRASHED"]
+WorkerStatus: TypeAlias = Literal["IDLE", "BUSY", "DRAINING", "SHUTDOWN", "CRASHED"]
+"""
+Status of worker in the worker pool
+"""
 
 
 class GetTestErrorByStatusCodeParams(TypedDict):
@@ -868,7 +929,7 @@ class PostGpuMemorysetByNameOrIdMemoryParams(TypedDict):
     name_or_id: str
 
 
-PostGpuMemorysetByNameOrIdMemoryRequest = list[LabeledMemoryInsert] | list[ScoredMemoryInsert]
+PostGpuMemorysetByNameOrIdMemoryRequest: TypeAlias = list[LabeledMemoryInsert] | list[ScoredMemoryInsert]
 
 
 class PatchGpuMemorysetByNameOrIdMemoriesParams(TypedDict):
@@ -1104,6 +1165,10 @@ class GetWorkerByWorkerIdParams(TypedDict):
 
 class GetTelemetryPredictionByPredictionIdParams(TypedDict):
     prediction_id: str
+    calc_neighborhood_density: NotRequired[bool]
+    """
+    Calculate neighborhood density
+    """
 
 
 class PatchTelemetryPredictionByPredictionIdParams(TypedDict):
@@ -1142,7 +1207,7 @@ class DeleteTelemetryFeedbackCategoryByNameOrIdParams(TypedDict):
     name_or_id: str
 
 
-PutTelemetryPredictionFeedbackRequest = list[PredictionFeedbackRequest]
+PutTelemetryPredictionFeedbackRequest: TypeAlias = list[PredictionFeedbackRequest]
 
 
 class GetAgentsBootstrapClassificationModelByJobIdParams(TypedDict):
@@ -1195,6 +1260,8 @@ class ClassificationMetrics(TypedDict):
     pr_auc: NotRequired[float | None]
     pr_curve: NotRequired[PRCurve | None]
     roc_curve: NotRequired[ROCCurve | None]
+    confusion_matrix: NotRequired[list[list[int]] | None]
+    warnings: NotRequired[list[str]]
 
 
 class ClassificationModelMetadata(TypedDict):
@@ -1418,7 +1485,7 @@ class HTTPValidationError(TypedDict):
 
 
 class InvalidInputErrorResponse(TypedDict):
-    status_code: NotRequired[int]
+    status_code: Literal[422]
     validation_issues: list[FieldValidationError]
 
 
@@ -1478,6 +1545,7 @@ class LabelPredictionWithMemoriesAndFeedback(TypedDict):
     tags: list[str]
     explanation: str | None
     memory_id: str | None
+    is_in_dense_neighborhood: NotRequired[bool | None]
     feedbacks: list[PredictionFeedback]
 
 
@@ -1549,10 +1617,10 @@ class ListPredictionsRequest(TypedDict):
     start_timestamp: NotRequired[str | None]
     end_timestamp: NotRequired[str | None]
     memory_id: NotRequired[str | None]
+    expected_label_match: NotRequired[bool | None]
     limit: NotRequired[int]
     offset: NotRequired[int | None]
     sort: NotRequired[PredictionSort]
-    expected_label_match: NotRequired[bool | None]
 
 
 class MemorysetAnalysisConfigs(TypedDict):
@@ -1631,10 +1699,10 @@ class WorkerInfo(TypedDict):
     config: dict[str, str | float | int | bool | dict[str, str] | None]
 
 
-PatchGpuMemorysetByNameOrIdMemoryRequest = LabeledMemoryUpdate | ScoredMemoryUpdate
+PatchGpuMemorysetByNameOrIdMemoryRequest: TypeAlias = LabeledMemoryUpdate | ScoredMemoryUpdate
 
 
-PatchGpuMemorysetByNameOrIdMemoriesRequest = list[LabeledMemoryUpdate] | list[ScoredMemoryUpdate]
+PatchGpuMemorysetByNameOrIdMemoriesRequest: TypeAlias = list[LabeledMemoryUpdate] | list[ScoredMemoryUpdate]
 
 
 class CascadingEditSuggestion(TypedDict):
@@ -1862,7 +1930,7 @@ class OrcaClient(Client):
         follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
         timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
         extensions: RequestExtensions | None = None,
-    ) -> bool:
+    ) -> Literal[True]:
         """Return true only when called with a valid root API key; otherwise 401 Unauthenticated."""
         pass
 
@@ -1896,7 +1964,7 @@ class OrcaClient(Client):
         follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
         timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
         extensions: RequestExtensions | None = None,
-    ) -> bool:
+    ) -> Literal[True]:
         """Returns true if the api key header is valid for the org (will be false for admin api key)"""
         pass
 

orca_sdk/credentials.py

@@ -1,10 +1,8 @@
-import os
 from datetime import datetime
-from typing import Literal, NamedTuple
+from typing import Literal
 
 import httpx
-from httpx import ConnectError, Headers, HTTPTransport
-from typing_extensions import deprecated
+from httpx import ConnectError, Headers
 
 from .async_client import OrcaAsyncClient
 from .client import OrcaClient
@@ -132,9 +130,6 @@ class OrcaCredentials:
         client = OrcaClient._resolve_client()
         client.DELETE("/auth/api_key/{name_or_id}", params={"name_or_id": name})
 
-    # TODO: remove deprecated methods after 2026-01-01
-
-    @deprecated("Use `OrcaClient.api_key` instead")
     @staticmethod
     def set_api_key(api_key: str, check_validity: bool = True):
         """
@@ -158,21 +153,25 @@ class OrcaCredentials:
         async_client = OrcaAsyncClient._resolve_client()
         async_client.api_key = api_key
 
-    @deprecated("Use `OrcaClient.base_url` instead")
     @staticmethod
     def get_api_url() -> str:
         """
         Get the base URL of the Orca API that is currently being used
         """
         client = OrcaClient._resolve_client()
+        async_client = OrcaAsyncClient._resolve_client()
+        if client.base_url != async_client.base_url:
+            raise RuntimeError("The base URL of the sync and async clients do not match")
         return str(client.base_url)
 
-    @deprecated("Use `OrcaClient.base_url` instead")
     @staticmethod
     def set_api_url(url: str, check_validity: bool = True):
         """
         Set the base URL for the Orca API
 
+        Note:
+            The base URL can also be provided by setting the `ORCA_API_URL` environment variable
+
         Args:
             url: The base URL to set
             check_validity: Whether to check if there is an API running at the given base URL
@@ -197,7 +196,6 @@ class OrcaCredentials:
         if check_validity:
             OrcaCredentials.is_healthy()
 
-    @deprecated("Use `OrcaClient.headers` instead")
     @staticmethod
     def set_api_headers(headers: dict[str, str]):
         """

orca_sdk/datasource.py

@@ -202,10 +202,10 @@ class Datasource:
             ValueError: If a datasource already exists and if_exists is `"error"`
         """
         if description is None or isinstance(description, str):
-            description = {dataset_name: description for dataset_name in dataset_dict.keys()}
+            description = {str(dataset_name): description for dataset_name in dataset_dict.keys()}
         return {
-            dataset_name: cls.from_hf_dataset(
-                f"{name}_{dataset_name}", dataset, if_exists=if_exists, description=description[dataset_name]
+            str(dataset_name): cls.from_hf_dataset(
+                f"{name}_{dataset_name}", dataset, if_exists=if_exists, description=description[str(dataset_name)]
             )
             for dataset_name, dataset in dataset_dict.items()
         }

orca_sdk/memoryset.py

@@ -32,13 +32,16 @@ from .client import (
     FilterItem,
 )
 from .client import LabeledMemory as LabeledMemoryResponse
-from .client import LabeledMemoryInsert
+from .client import (
+    LabeledMemoryInsert,
+)
 from .client import LabeledMemoryLookup as LabeledMemoryLookupResponse
 from .client import (
     LabeledMemoryUpdate,
     LabeledMemoryWithFeedbackMetrics,
     LabelPredictionMemoryLookup,
     LabelPredictionWithMemoriesAndFeedback,
+    ListPredictionsRequest,
     MemoryMetrics,
     MemorysetAnalysisConfigs,
     MemorysetMetadata,
@@ -46,16 +49,18 @@ from .client import (
     MemorysetUpdate,
     MemoryType,
     OrcaClient,
-    PredictionFeedback,
 )
 from .client import ScoredMemory as ScoredMemoryResponse
-from .client import ScoredMemoryInsert
+from .client import (
+    ScoredMemoryInsert,
+)
 from .client import ScoredMemoryLookup as ScoredMemoryLookupResponse
 from .client import (
     ScoredMemoryUpdate,
     ScoredMemoryWithFeedbackMetrics,
     ScorePredictionMemoryLookup,
     ScorePredictionWithMemoriesAndFeedback,
+    TelemetryField,
     TelemetryFilterItem,
     TelemetrySortOptions,
 )
@@ -157,9 +162,10 @@ def _parse_filter_item_from_tuple(input: FilterItemTuple) -> FilterItem | Teleme
             raise ValueError("Like filters are not supported on metric columns")
         op = cast(Literal["==", "!=", ">", ">=", "<", "<=", "in", "not in"], op)
         value = cast(float | int | list[float] | list[int], value)
-        return TelemetryFilterItem(field=field, op=op, value=value)
+        return TelemetryFilterItem(field=cast(TelemetryField, tuple(field)), op=op, value=value)
 
-    return FilterItem(field=field, op=op, value=value)
+    # Convert list to tuple for FilterItem field type
+    return FilterItem(field=tuple(field), op=op, value=value)  # type: ignore[assignment]
 
 
 def _parse_sort_item_from_tuple(
@@ -183,7 +189,8 @@ def _parse_sort_item_from_tuple(
             raise ValueError("Lookup must follow the format `lookup.count`")
         if field[1] != "count":
             raise ValueError("Lookup can only be sorted on count")
-    return TelemetrySortOptions(field=field, direction=input[1])
+    # Convert list to tuple for TelemetryField type
+    return TelemetrySortOptions(field=cast(TelemetryField, tuple(field)), direction=input[1])
 
 
 def _parse_memory_insert(memory: dict[str, Any], type: MemoryType) -> LabeledMemoryInsert | ScoredMemoryInsert:
@@ -593,16 +600,18 @@ class LabeledMemory(MemoryBase):
         """
 
         client = OrcaClient._resolve_client()
+        request_json: ListPredictionsRequest = {
+            "memory_id": self.memory_id,
+            "limit": limit,
+            "offset": offset,
+            "tag": tag,
+            "expected_label_match": expected_label_match,
+        }
+        if sort:
+            request_json["sort"] = sort
         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,
-            },
+            json=request_json,
         )
 
         # Filter to only classification predictions and convert to ClassificationPrediction objects
@@ -808,16 +817,18 @@ class ScoredMemory(MemoryBase):
             List of RegressionPrediction objects that used this memory
         """
         client = OrcaClient._resolve_client()
+        request_json: ListPredictionsRequest = {
+            "memory_id": self.memory_id,
+            "limit": limit,
+            "offset": offset,
+            "tag": tag,
+            "expected_label_match": expected_label_match,
+        }
+        if sort:
+            request_json["sort"] = sort
         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,
-            },
+            json=request_json,
         )
 
         # Filter to only regression predictions and convert to RegressionPrediction objects
@@ -940,8 +951,6 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
     index_params: dict[str, Any]
     hidden: bool
 
-    _batch_size = 32  # max number of memories to insert/update/delete in a single API call
-
     def __init__(self, metadata: MemorysetMetadata):
         # for internal use only, do not document
         if metadata["pretrained_embedding_model_name"]:
@@ -2532,7 +2541,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             ]
         )
 
-    def insert(self, items: Iterable[dict[str, Any]] | dict[str, Any]) -> None:
+    def insert(self, items: Iterable[dict[str, Any]] | dict[str, Any], *, batch_size: int = 32) -> None:
         """
         Insert memories into the memoryset
 
@@ -2546,17 +2555,21 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 - `source_id`: Optional unique ID of the memory in a system of reference
                 - `...`: Any other metadata to store for the memory
 
+            batch_size: Number of memories to insert in a single API call
+
         Examples:
             >>> memoryset.insert([
             ...     {"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"},
             ... ])
         """
+        if batch_size <= 0 or batch_size > 500:
+            raise ValueError("batch_size must be between 1 and 500")
         client = OrcaClient._resolve_client()
         items = cast(list[dict[str, Any]], [items]) if isinstance(items, dict) else list(items)
         # insert memories in batches to avoid API timeouts
-        for i in range(0, len(items), self._batch_size):
-            batch = items[i : i + self._batch_size]
+        for i in range(0, len(items), batch_size):
+            batch = items[i : i + batch_size]
             client.POST(
                 "/gpu/memoryset/{name_or_id}/memory",
                 params={"name_or_id": self.id},
@@ -2568,7 +2581,7 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
 
         self.refresh()
 
-    async def ainsert(self, items: Iterable[dict[str, Any]] | dict[str, Any]) -> None:
+    async def ainsert(self, items: Iterable[dict[str, Any]] | dict[str, Any], *, batch_size: int = 32) -> None:
         """
         Asynchronously insert memories into the memoryset
 
@@ -2583,17 +2596,21 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 - `partition_id`: Optional partition ID of the memory
                 - `...`: Any other metadata to store for the memory
 
+            batch_size: Number of memories to insert in a single API call
+
         Examples:
             >>> await memoryset.ainsert([
             ...     {"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"},
             ... ])
         """
+        if batch_size <= 0 or batch_size > 500:
+            raise ValueError("batch_size must be between 1 and 500")
         client = OrcaAsyncClient._resolve_client()
         items = cast(list[dict[str, Any]], [items]) if isinstance(items, dict) else list(items)
         # insert memories in batches to avoid API timeouts
-        for i in range(0, len(items), self._batch_size):
-            batch = items[i : i + self._batch_size]
+        for i in range(0, len(items), batch_size):
+            batch = items[i : i + batch_size]
             await client.POST(
                 "/gpu/memoryset/{name_or_id}/memory",
                 params={"name_or_id": self.id},
@@ -2682,14 +2699,16 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             ]
 
     @overload
-    def update(self, updates: dict[str, Any]) -> MemoryT:
+    def update(self, updates: dict[str, Any], *, batch_size: int = 32) -> MemoryT:
         pass
 
     @overload
-    def update(self, updates: Iterable[dict[str, Any]]) -> list[MemoryT]:
+    def update(self, updates: Iterable[dict[str, Any]], *, batch_size: int = 32) -> list[MemoryT]:
         pass
 
-    def update(self, updates: dict[str, Any] | Iterable[dict[str, Any]]) -> MemoryT | list[MemoryT]:
+    def update(
+        self, updates: dict[str, Any] | Iterable[dict[str, Any]], *, batch_size: int = 32
+    ) -> MemoryT | list[MemoryT]:
         """
         Update one or multiple memories in the memoryset
 
@@ -2704,6 +2723,8 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
                 - `partition_id`: Optional new partition ID of the memory
                 - `...`: Optional new values for metadata properties
 
+            batch_size: Number of memories to update in a single API call
+
         Returns:
             Updated memory or list of updated memories
 
@@ -2722,12 +2743,14 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             ...     for m in memoryset.query(filters=[("tag", "==", "happy")])
             ... )
         """
+        if batch_size <= 0 or batch_size > 500:
+            raise ValueError("batch_size must be between 1 and 500")
         client = OrcaClient._resolve_client()
         updates_list = cast(list[dict[str, Any]], [updates]) if isinstance(updates, dict) else list(updates)
         # update memories in batches to avoid API timeouts
         updated_memories: list[MemoryT] = []
-        for i in range(0, len(updates_list), self._batch_size):
-            batch = updates_list[i : i + self._batch_size]
+        for i in range(0, len(updates_list), batch_size):
+            batch = updates_list[i : i + batch_size]
             response = client.PATCH(
                 "/gpu/memoryset/{name_or_id}/memories",
                 params={"name_or_id": self.id},
@@ -2803,12 +2826,13 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             },
         )
 
-    def delete(self, memory_id: str | Iterable[str]) -> None:
+    def delete(self, memory_id: str | Iterable[str], *, batch_size: int = 32) -> None:
         """
         Delete memories from the memoryset
 
         Params:
             memory_id: unique identifiers of the memories to delete
+            batch_size: Number of memories to delete in a single API call
 
         Examples:
             Delete a single memory:
@@ -2821,11 +2845,13 @@ class MemorysetBase(Generic[MemoryT, MemoryLookupT], ABC):
             ... )
 
         """
+        if batch_size <= 0 or batch_size > 500:
+            raise ValueError("batch_size must be between 1 and 500")
         client = OrcaClient._resolve_client()
         memory_ids = [memory_id] if isinstance(memory_id, str) else list(memory_id)
         # delete memories in batches to avoid API timeouts
-        for i in range(0, len(memory_ids), self._batch_size):
-            batch = memory_ids[i : i + self._batch_size]
+        for i in range(0, len(memory_ids), batch_size):
+            batch = memory_ids[i : i + batch_size]
             client.POST(
                 "/memoryset/{name_or_id}/memories/delete", params={"name_or_id": self.id}, json={"memory_ids": batch}
             )