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

orca_sdk/memoryset_test.py

@@ -856,7 +856,8 @@ def test_insert_memories(writable_memoryset: LabeledMemoryset):
         [
             dict(value="tomato soup is my favorite", label=0),
             dict(value="cats are fun to play with", label=1),
-        ]
+        ],
+        batch_size=1,
     )
     writable_memoryset.refresh()
     assert writable_memoryset.length == prev_length + 2
@@ -897,7 +898,8 @@ def test_update_memories(writable_memoryset: LabeledMemoryset, hf_dataset: Datas
         [
             dict(memory_id=memory_ids[0], value="i love soup so much"),
             dict(memory_id=memory_ids[1], value="cats are so cute"),
-        ]
+        ],
+        batch_size=1,
     )
     assert updated_memories[0].value == "i love soup so much"
     assert updated_memories[1].value == "cats are so cute"
@@ -916,7 +918,7 @@ def test_delete_memories(writable_memoryset: LabeledMemoryset):
 
     # test deleting multiple memories
     prev_length = writable_memoryset.length
-    writable_memoryset.delete([writable_memoryset[0].memory_id, writable_memoryset[1].memory_id])
+    writable_memoryset.delete([writable_memoryset[0].memory_id, writable_memoryset[1].memory_id], batch_size=1)
     assert writable_memoryset.length == prev_length - 2
 
 

orca_sdk/regression_model.py

@@ -10,6 +10,7 @@ from datasets import Dataset
 from ._shared.metrics import RegressionMetrics, calculate_regression_metrics
 from ._utils.common import UNSET, CreateMode, DropMode
 from .client import (
+    ListPredictionsRequest,
     OrcaClient,
     PredictiveModelUpdate,
     RARHeadType,
@@ -295,6 +296,7 @@ class RegressionModel:
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
         use_gpu: bool = True,
+        batch_size: int = 100,
     ) -> RegressionPrediction: ...
 
     @overload
@@ -313,6 +315,7 @@ class RegressionModel:
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
         use_gpu: bool = True,
+        batch_size: int = 100,
     ) -> list[RegressionPrediction]: ...
 
     # TODO: add filter support
@@ -331,6 +334,7 @@ class RegressionModel:
             "ignore_partitions", "include_global", "exclude_global", "only_global"
         ] = "include_global",
         use_gpu: bool = True,
+        batch_size: int = 100,
     ) -> RegressionPrediction | list[RegressionPrediction]:
         """
         Make predictions using the regression model.
@@ -355,6 +359,7 @@ class RegressionModel:
                 * `"exclude_global"`: Exclude global memories
                 * `"only_global"`: Only include global memories
             use_gpu: Whether to use GPU for the prediction (defaults to True)
+            batch_size: Number of values to process in a single API call
 
         Returns:
             Single RegressionPrediction or list of RegressionPrediction objects
@@ -366,6 +371,8 @@ class RegressionModel:
         """
         if timeout_seconds <= 0:
             raise ValueError("timeout_seconds must be a positive integer")
+        if batch_size <= 0 or batch_size > 500:
+            raise ValueError("batch_size must be between 1 and 500")
 
         if use_gpu:
             endpoint = "/gpu/regression_model/{name_or_id}/prediction"
@@ -374,75 +381,98 @@ class RegressionModel:
 
         telemetry_on, telemetry_sync = _get_telemetry_config(save_telemetry)
         client = OrcaClient._resolve_client()
-        request_json: RegressionPredictionRequest = {
-            "input_values": value if isinstance(value, list) else [value],
-            "memoryset_override_name_or_id": self._memoryset_override_id,
-            "expected_scores": (
-                expected_scores
-                if isinstance(expected_scores, list)
-                else [expected_scores] if expected_scores is not None else None
-            ),
-            "tags": list(tags or set()),
-            "save_telemetry": telemetry_on,
-            "save_telemetry_synchronously": telemetry_sync,
-            "prompt": prompt,
-            "use_lookup_cache": use_lookup_cache,
-            "ignore_unlabeled": ignore_unlabeled,
-            "partition_filter_mode": partition_filter_mode,
-        }
-        # Don't send partition_ids when partition_filter_mode is "ignore_partitions"
-        if partition_filter_mode != "ignore_partitions":
-            request_json["partition_ids"] = partition_id
-        response = client.POST(
-            endpoint,
-            params={"name_or_id": self.id},
-            json=request_json,
-            timeout=timeout_seconds,
-        )
 
-        if telemetry_on and any(p["prediction_id"] is None for p in response):
-            raise RuntimeError("Failed to save prediction to database.")
+        # Convert to list for batching
+        values = value if isinstance(value, list) else [value]
+        if isinstance(expected_scores, list) and len(expected_scores) != len(values):
+            raise ValueError("Invalid input: \n\texpected_scores must be the same length as values")
+        if isinstance(partition_id, list) and len(partition_id) != len(values):
+            raise ValueError("Invalid input: \n\tpartition_id must be the same length as values")
+
+        if isinstance(expected_scores, list):
+            expected_scores = expected_scores
+        elif expected_scores is not None:
+            expected_scores = [float(expected_scores)] * len(values)
+
+        predictions: list[RegressionPrediction] = []
+        for i in range(0, len(values), batch_size):
+            batch_values = values[i : i + batch_size]
+            batch_expected_scores = expected_scores[i : i + batch_size] if expected_scores else None
+
+            request_json: RegressionPredictionRequest = {
+                "input_values": batch_values,
+                "memoryset_override_name_or_id": self._memoryset_override_id,
+                "expected_scores": batch_expected_scores,
+                "tags": list(tags or set()),
+                "save_telemetry": telemetry_on,
+                "save_telemetry_synchronously": telemetry_sync,
+                "prompt": prompt,
+                "use_lookup_cache": use_lookup_cache,
+                "ignore_unlabeled": ignore_unlabeled,
+                "partition_filter_mode": partition_filter_mode,
+            }
+            if partition_filter_mode != "ignore_partitions":
+                request_json["partition_ids"] = (
+                    partition_id[i : i + batch_size] if isinstance(partition_id, list) else partition_id
+                )
+
+            response = client.POST(
+                endpoint,
+                params={"name_or_id": self.id},
+                json=request_json,
+                timeout=timeout_seconds,
+            )
 
-        predictions = [
-            RegressionPrediction(
-                prediction_id=prediction["prediction_id"],
-                label=None,
-                label_name=None,
-                score=prediction["score"],
-                confidence=prediction["confidence"],
-                anomaly_score=prediction["anomaly_score"],
-                memoryset=self.memoryset,
-                model=self,
-                logits=None,
-                input_value=input_value,
+            if telemetry_on and any(p["prediction_id"] is None for p in response):
+                raise RuntimeError("Failed to save prediction to database.")
+
+            predictions.extend(
+                RegressionPrediction(
+                    prediction_id=prediction["prediction_id"],
+                    label=None,
+                    label_name=None,
+                    score=prediction["score"],
+                    confidence=prediction["confidence"],
+                    anomaly_score=prediction["anomaly_score"],
+                    memoryset=self.memoryset,
+                    model=self,
+                    logits=None,
+                    input_value=input_value,
+                )
+                for prediction, input_value in zip(response, batch_values)
             )
-            for prediction, input_value in zip(response, value if isinstance(value, list) else [value])
-        ]
+
         self._last_prediction_was_batch = isinstance(value, list)
         self._last_prediction = predictions[-1]
         return predictions if isinstance(value, list) else predictions[0]
 
     def predictions(
         self,
-        limit: int = 100,
+        limit: int | None = None,
         offset: int = 0,
         tag: str | None = None,
         sort: list[tuple[Literal["anomaly_score", "confidence", "timestamp"], Literal["asc", "desc"]]] = [],
+        batch_size: int = 100,
     ) -> list[RegressionPrediction]:
         """
         Get a list of predictions made by this model
 
         Params:
-            limit: Optional maximum number of predictions to return
+            limit: Maximum number of predictions to return. If `None`, returns all predictions
+                by automatically paginating through results.
             offset: Optional offset of the first prediction to return
             tag: Optional tag to filter predictions by
             sort: Optional list of columns and directions to sort the predictions by.
                 Predictions can be sorted by `created_at`, `confidence`, `anomaly_score`, or `score`.
+            batch_size: Number of predictions to fetch in a single API call
 
         Returns:
             List of score predictions
 
         Examples:
+            Get all predictions with a specific tag:
+            >>> predictions = model.predictions(tag="evaluation")
+
             Get the last 3 predictions:
             >>> predictions = model.predictions(limit=3, sort=[("created_at", "desc")])
             [
@@ -455,34 +485,61 @@ class RegressionModel:
             >>> predictions = model.predictions(sort=[("confidence", "desc")], offset=1, limit=1)
             [RegressionPrediction({score: 4.2, confidence: 0.90, anomaly_score: 0.1, input_value: 'Good service'})]
         """
+        if batch_size <= 0 or batch_size > 500:
+            raise ValueError("batch_size must be between 1 and 500")
+        if limit == 0:
+            return []
+
         client = OrcaClient._resolve_client()
-        predictions = client.POST(
-            "/telemetry/prediction",
-            json={
+        all_predictions: list[RegressionPrediction] = []
+
+        if limit is not None and limit < batch_size:
+            pages = [(offset, limit)]
+        else:
+            # automatically paginate the requests if necessary
+            total = client.POST(
+                "/telemetry/prediction/count",
+                json={
+                    "model_id": self.id,
+                    "tag": tag,
+                },
+            )
+            max_limit = max(total - offset, 0)
+            limit = min(limit, max_limit) if limit is not None else max_limit
+            pages = [(o, min(batch_size, limit - (o - offset))) for o in range(offset, offset + limit, batch_size)]
+
+        for current_offset, current_limit in pages:
+            request_json: ListPredictionsRequest = {
                 "model_id": self.id,
-                "limit": limit,
-                "offset": offset,
-                "sort": [list(sort_item) for sort_item in sort],
+                "limit": current_limit,
+                "offset": current_offset,
                 "tag": tag,
-            },
-        )
-        return [
-            RegressionPrediction(
-                prediction_id=prediction["prediction_id"],
-                label=None,
-                label_name=None,
-                score=prediction["score"],
-                confidence=prediction["confidence"],
-                anomaly_score=prediction["anomaly_score"],
-                memoryset=self.memoryset,
-                model=self,
-                telemetry=prediction,
-                logits=None,
-                input_value=None,
+            }
+            if sort:
+                request_json["sort"] = sort
+            response = client.POST(
+                "/telemetry/prediction",
+                json=request_json,
             )
-            for prediction in predictions
-            if "score" in prediction
-        ]
+            all_predictions.extend(
+                RegressionPrediction(
+                    prediction_id=prediction["prediction_id"],
+                    label=None,
+                    label_name=None,
+                    score=prediction["score"],
+                    confidence=prediction["confidence"],
+                    anomaly_score=prediction["anomaly_score"],
+                    memoryset=self.memoryset,
+                    model=self,
+                    telemetry=prediction,
+                    logits=None,
+                    input_value=None,
+                )
+                for prediction in response
+                if "score" in prediction
+            )
+
+        return all_predictions
 
     def _evaluate_datasource(
         self,

orca_sdk/regression_model_test.py

@@ -183,14 +183,18 @@ def test_evaluate_dataset_with_nones_raises_error(regression_model: RegressionMo
 
 
 def test_evaluate_with_telemetry(regression_model, eval_dataset: Dataset):
-    result = regression_model.evaluate(eval_dataset, record_predictions=True, tags={"test"})
+    result = regression_model.evaluate(eval_dataset, record_predictions=True, tags={"test"}, batch_size=2)
     assert result is not None
     assert isinstance(result, RegressionMetrics)
-    predictions = regression_model.predictions(tag="test")
+    predictions = regression_model.predictions(tag="test", batch_size=100, sort=[("timestamp", "asc")])
     assert len(predictions) == 4
     assert all(p.tags == {"test"} for p in predictions)
     assert all(p.expected_score is not None for p in predictions)
-    assert all(np.allclose(p.expected_score, s) for p, s in zip(predictions, eval_dataset["score"]))
+    prediction_expected_scores = [p.expected_score for p in predictions]
+    eval_expected_scores = list(eval_dataset["score"])
+    assert all(
+        np.allclose(p, s, atol=1e-3) for p, s in zip(prediction_expected_scores, eval_expected_scores)
+    ), f"Prediction expected scores: {prediction_expected_scores} do not match eval expected scores: {eval_expected_scores}"
 
 
 def test_evaluate_with_partition_column_dataset(partitioned_regression_model: RegressionModel):
@@ -328,7 +332,7 @@ def test_evaluate_with_partition_column_datasource(partitioned_regression_model:
 
 
 def test_predict(regression_model: RegressionModel):
-    predictions = regression_model.predict(["Do you love soup?", "Are cats cute?"])
+    predictions = regression_model.predict(["Do you love soup?", "Are cats cute?"], batch_size=1)
     assert len(predictions) == 2
     assert predictions[0].prediction_id is not None
     assert predictions[1].prediction_id is not None

{orca_sdk-0.1.8.dist-info → orca_sdk-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orca_sdk
-Version: 0.1.8
+Version: 0.1.10
 Summary: SDK for interacting with Orca Services
 License-Expression: Apache-2.0
 Author: Orca DB Inc.
@@ -11,13 +11,13 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
-Requires-Dist: datasets (>=3.1.0,<4)
-Requires-Dist: gradio (>=5.44.1,<6)
+Requires-Dist: datasets (>=4.4.0,<5)
+Requires-Dist: gradio (>=6.0.0,<7)
 Requires-Dist: httpx (>=0.28.1)
 Requires-Dist: httpx-retries (>=0.4.3,<0.5.0)
 Requires-Dist: numpy (>=2.1.0,<3)
 Requires-Dist: pandas (>=2.2.3,<3)
-Requires-Dist: pyarrow (>=18.0.0,<19)
+Requires-Dist: pyarrow (>=22.0.0,<23)
 Requires-Dist: python-dotenv (>=1.1.0)
 Requires-Dist: scikit-learn (>=1.6.1,<2)
 Requires-Dist: torch (>=2.8.0,<3)

{orca_sdk-0.1.8.dist-info → orca_sdk-0.1.10.dist-info}/RECORD

@@ -3,13 +3,13 @@ orca_sdk/_shared/__init__.py,sha256=3Kt0Hu3QLI5FEp9nqGTxqAm3hAoBJKcagfaGQZ-lbJQ,
 orca_sdk/_shared/metrics.py,sha256=faeL1B1ftmns1ikfKrIlU3xOn6j0iAGLNUupxvAFza8,24968
 orca_sdk/_shared/metrics_test.py,sha256=vDIXoj8EuuLcdPJz_7EiVPgQ-FXiVT81JG30jxsg9HM,20752
 orca_sdk/_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-orca_sdk/_utils/analysis_ui.py,sha256=nT-M_YcNRCVPQzvuqYNFKnNHhYkADYBvq1GlIUePrWw,9232
+orca_sdk/_utils/analysis_ui.py,sha256=R0xc4RyJKyBHJEEF_ztI4Dm5w8Y1uF0Wpkn4LQgXqBE,9258
 orca_sdk/_utils/analysis_ui_style.css,sha256=q_ba_-_KtgztepHg829zLzypaxKayl7ySC1-oYDzV3k,836
 orca_sdk/_utils/auth.py,sha256=nC252O171_3_wn4KBAN7kg8GNvoZFiQ5Xtzkrm5dWDo,2645
 orca_sdk/_utils/auth_test.py,sha256=ygVWv1Ex53LaxIP7p2hzPHl8l9qYyBD5IGmEFJMps6s,1056
 orca_sdk/_utils/common.py,sha256=wUm2pNDWytEecC5WiDWd02-yCZw3Akx0bIutG4lHsFA,805
-orca_sdk/_utils/data_parsing.py,sha256=gkAwWEC8qRt3vRUObe7n7Pr0azOayNwc2yFY04WFp7E,5220
-orca_sdk/_utils/data_parsing_test.py,sha256=mJIPDOfs9fTsrQO_oMGonVAtNCA3_gcw_3qSmOX5Q4g,8978
+orca_sdk/_utils/data_parsing.py,sha256=5vaTpvUOS-ldlcgnSARYw7s9mce-imzkU7kA48-pdIM,5396
+orca_sdk/_utils/data_parsing_test.py,sha256=u7BEjxtsU9gMs3tAZI0lJ--vOLlwKwH3hemdCedzxA0,8826
 orca_sdk/_utils/pagination.py,sha256=986z0QPZixrZeurJWorF6eMgnTRdDF84AagEA6qNbMw,4245
 orca_sdk/_utils/pagination_test.py,sha256=BUylCrcHnwoKEBmMUzVr0lwLpA35ivcCwdBK4rMw9y8,4887
 orca_sdk/_utils/prediction_result_ui.css,sha256=sqBlkRLnovb5X5EcUDdB6iGpH63nVRlTW4uAmXuD0WM,258
@@ -17,25 +17,25 @@ orca_sdk/_utils/prediction_result_ui.py,sha256=Ur_FY7dz3oWNmtPiP3Wl3yRlEMgK8q9Uf
 orca_sdk/_utils/tqdm_file_reader.py,sha256=Lw7Cg1UgNuRUoN6jjqZb-IlV00H-kbRcrZLdudr1GxE,324
 orca_sdk/_utils/value_parser.py,sha256=c3qMABCCDQcIjn9N1orYYnlRwDW9JWdGwW_2TDZPLdI,1286
 orca_sdk/_utils/value_parser_test.py,sha256=OybsiC-Obi32RRi9NIuwrVBRAnlyPMV1xVAaevSrb7M,1079
-orca_sdk/async_client.py,sha256=U7B8yLbeB5STPr3dkeTTwPkDHou2nBCJ9CLBPulrXVY,134893
-orca_sdk/classification_model.py,sha256=jy8h5Kb4cesFqm1sat4JuxAs9uH53rUvOrlqKID-mwA,46294
-orca_sdk/classification_model_test.py,sha256=vBn7KBb9-ACuJEdzW50n54Fn6Mh9iEYbn1197lE8-yI,36997
-orca_sdk/client.py,sha256=9AaEYrfK7HZEpsSSih8mRuTUiaQ4yk2n0jU92oLusGY,133942
+orca_sdk/async_client.py,sha256=PM7N-ggmtucfcUF1vQGtTZOCJpSNTOgd7l3LDNF5kP4,137192
+orca_sdk/classification_model.py,sha256=C58euWnNvwXnthR9RtVVCOcgPEbxCjjp3sHMb86V6YA,50197
+orca_sdk/classification_model_test.py,sha256=ElqxtR6gNwwk8dNXwfwAhpT7l0ZIP3H4pHmOyFXyTWk,37370
+orca_sdk/client.py,sha256=SKZv3zGG6OwLe_FlB5wL2cxltOLPCcHvoo2CbMwyKgA,136241
 orca_sdk/conftest.py,sha256=0O1VY-SPKNAvi9fBLdY1RMnYVgZvMjP92y99bNAqqiw,12461
-orca_sdk/credentials.py,sha256=80_1r8n5jruEvN_E629SaRrRhKvF_NhWUEZyZzPXkqQ,6620
+orca_sdk/credentials.py,sha256=2SwC3tq5akP-F_u2s4xMZDp8mlsKMUT1T5T9Z99-eSY,6588
 orca_sdk/credentials_test.py,sha256=TLbXJMz3IlThvtSrHeLM7jRsKnrncA_ahOTpHg15Ei4,4089
-orca_sdk/datasource.py,sha256=6QaccghiyFEUSFcqnwjIJzpgIh9Id0snJk2EqViqPsU,22356
+orca_sdk/datasource.py,sha256=Qn5QloE84UXeyPk2wcy1lWe5wmh1iDBS044eWnxck_E,22371
 orca_sdk/datasource_test.py,sha256=sCk3IcQJbDut5oN4Wf7PXhTxyMwalxMuCXJekSxy9wk,16665
 orca_sdk/embedding_model.py,sha256=vLGnlO9I-cN1lklNBl_LxZ8m9oK3vkegFOpvYYw8u8g,28038
 orca_sdk/embedding_model_test.py,sha256=Lc6fZ0ifT0hh6ldkUfjwMPcP6OgN0Umlzu8XDLs7UO4,8144
 orca_sdk/job.py,sha256=wHwVt-s7i-v8udhLGybB-90Kp4dwOLrY806bE4Tam5Q,13092
 orca_sdk/job_test.py,sha256=nRSWxd_1UIfrj9oMVvrXjt6OBkBpddYAjb2y6P-DTUg,4327
-orca_sdk/memoryset.py,sha256=-bTQJLzr8Fdt3LIaCEliO2jT_LxcaoaqxETY42vH0mA,156070
-orca_sdk/memoryset_test.py,sha256=X-gAOz-BA4ByXAitI19GQsktnL2xRGcLus8AHj5mL48,51218
-orca_sdk/regression_model.py,sha256=Ga0TZLm1Qe_bQXoP90BrMiyvZ8bsJ33Z9I6RgT86t2c,31339
-orca_sdk/regression_model_test.py,sha256=DfWLkqxB835jjwM-sj1uxQ6Yz_ZBMnt8EHjdfnHsRnU,25103
+orca_sdk/memoryset.py,sha256=06v34fHabpkEaOv9VCKc0NhpMi_mNZGbQP_9GiW_nuE,157157
+orca_sdk/memoryset_test.py,sha256=O2o42XETtffXtZy0kbLk2b8cUDXU-w2ZAzXLi5-vDPQ,51278
+orca_sdk/regression_model.py,sha256=AXRzJG15sDJQSiDCDfRdcLnZDNkJWORYjhHqKyyL-Fc,33960
+orca_sdk/regression_model_test.py,sha256=90EyrhaMk1kTf87RFkMNz1PTItmeUISs6AvHmyp08DU,25447
 orca_sdk/telemetry.py,sha256=ZyCMiyyo_SchjadWZH55TlLrC4Ucq5S316NbW26LL4Y,27834
 orca_sdk/telemetry_test.py,sha256=eT66C5lFdNg-pQdo2I__BP7Tn5fTc9aTkVo9ZhWwhU0,5519
-orca_sdk-0.1.8.dist-info/METADATA,sha256=EIDzSzxtf7D5LEuAdakH5rIy1XHqeEC1HhqsrF_oT5c,3710
-orca_sdk-0.1.8.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-orca_sdk-0.1.8.dist-info/RECORD,,
+orca_sdk-0.1.10.dist-info/METADATA,sha256=j_TIalbL8oztP39lnXjyAI6Aosvb6rnJKUc3gcuxD0k,3710
+orca_sdk-0.1.10.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+orca_sdk-0.1.10.dist-info/RECORD,,