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

orca_sdk/classification_model.py

@@ -9,13 +9,16 @@ from datasets import Dataset
 
 from ._shared.metrics import ClassificationMetrics, calculate_classification_metrics
 from ._utils.common import UNSET, CreateMode, DropMode
+from .async_client import OrcaAsyncClient
 from .client import (
     BootstrapClassificationModelMeta,
     BootstrapClassificationModelResult,
+    ClassificationEvaluationRequest,
     ClassificationModelMetadata,
+    OrcaClient,
+    PostClassificationModelByModelNameOrIdEvaluationParams,
     PredictiveModelUpdate,
     RACHeadType,
-    orca_api,
 )
 from .datasource import Datasource
 from .job import Job
@@ -199,7 +202,12 @@ class ClassificationModel:
                 raise ValueError(f"Model with name {name} already exists")
             elif if_exists == "open":
                 existing = cls.open(name)
-                for attribute in {"head_type", "memory_lookup_count", "num_classes", "min_memory_weight"}:
+                for attribute in {
+                    "head_type",
+                    "memory_lookup_count",
+                    "num_classes",
+                    "min_memory_weight",
+                }:
                     local_attribute = locals()[attribute]
                     existing_attribute = getattr(existing, attribute)
                     if local_attribute is not None and local_attribute != existing_attribute:
@@ -211,7 +219,8 @@ class ClassificationModel:
 
                 return existing
 
-        metadata = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        metadata = client.POST(
             "/classification_model",
             json={
                 "name": name,
@@ -240,7 +249,8 @@ class ClassificationModel:
         Raises:
             LookupError: If the classification model does not exist
         """
-        return cls(orca_api.GET("/classification_model/{name_or_id}", params={"name_or_id": name}))
+        client = OrcaClient._resolve_client()
+        return cls(client.GET("/classification_model/{name_or_id}", params={"name_or_id": name}))
 
     @classmethod
     def exists(cls, name_or_id: str) -> bool:
@@ -267,7 +277,8 @@ class ClassificationModel:
         Returns:
             List of handles to all classification models in the OrcaCloud
         """
-        return [cls(metadata) for metadata in orca_api.GET("/classification_model")]
+        client = OrcaClient._resolve_client()
+        return [cls(metadata) for metadata in client.GET("/classification_model")]
 
     @classmethod
     def drop(cls, name_or_id: str, if_not_exists: DropMode = "error"):
@@ -286,7 +297,8 @@ class ClassificationModel:
             LookupError: If the classification model does not exist and if_not_exists is `"error"`
         """
         try:
-            orca_api.DELETE("/classification_model/{name_or_id}", params={"name_or_id": name_or_id})
+            client = OrcaClient._resolve_client()
+            client.DELETE("/classification_model/{name_or_id}", params={"name_or_id": name_or_id})
             logging.info(f"Deleted model {name_or_id}")
         except LookupError:
             if if_not_exists == "error":
@@ -322,7 +334,8 @@ class ClassificationModel:
             update["description"] = description
         if locked is not UNSET:
             update["locked"] = locked
-        orca_api.PATCH("/classification_model/{name_or_id}", params={"name_or_id": self.id}, json=update)
+        client = OrcaClient._resolve_client()
+        client.PATCH("/classification_model/{name_or_id}", params={"name_or_id": self.id}, json=update)
         self.refresh()
 
     def lock(self) -> None:
@@ -344,6 +357,8 @@ class ClassificationModel:
         prompt: str | None = None,
         use_lookup_cache: bool = True,
         timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+        use_gpu: bool = True,
     ) -> list[ClassificationPrediction]:
         pass
 
@@ -358,6 +373,8 @@ class ClassificationModel:
         prompt: str | None = None,
         use_lookup_cache: bool = True,
         timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+        use_gpu: bool = True,
     ) -> ClassificationPrediction:
         pass
 
@@ -371,6 +388,8 @@ class ClassificationModel:
         prompt: str | None = None,
         use_lookup_cache: bool = True,
         timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+        use_gpu: bool = True,
     ) -> list[ClassificationPrediction] | ClassificationPrediction:
         """
         Predict label(s) for the given input value(s) grounded in similar memories
@@ -389,6 +408,9 @@ class ClassificationModel:
             prompt: Optional prompt to use for instruction-tuned embedding models
             use_lookup_cache: Whether to use cached lookup results for faster predictions
             timeout_seconds: Timeout in seconds for the request, defaults to 10 seconds
+            ignore_unlabeled: If True, only use labeled memories during lookup.
+                If False (default), allow unlabeled memories when necessary.
+            use_gpu: Whether to use GPU for the prediction (defaults to True)
 
         Returns:
             Label prediction or list of label predictions
@@ -421,6 +443,159 @@ class ClassificationModel:
             _parse_filter_item_from_tuple(filter) if isinstance(filter, tuple) else filter for filter in filters
         ]
 
+        if any(_is_metric_column(filter[0]) for filter in filters):
+            raise ValueError(f"Cannot filter on {filters} - telemetry filters are not supported for predictions")
+
+        if isinstance(expected_labels, int):
+            expected_labels = [expected_labels]
+        elif isinstance(expected_labels, str):
+            expected_labels = [self.memoryset.label_names.index(expected_labels)]
+        elif isinstance(expected_labels, list):
+            expected_labels = [
+                self.memoryset.label_names.index(label) if isinstance(label, str) else label
+                for label in expected_labels
+            ]
+
+        if use_gpu:
+            endpoint = "/gpu/classification_model/{name_or_id}/prediction"
+        else:
+            endpoint = "/classification_model/{name_or_id}/prediction"
+
+        telemetry_on, telemetry_sync = _get_telemetry_config(save_telemetry)
+        client = OrcaClient._resolve_client()
+        response = client.POST(
+            endpoint,
+            params={"name_or_id": self.id},
+            json={
+                "input_values": value if isinstance(value, list) else [value],
+                "memoryset_override_name_or_id": self._memoryset_override_id,
+                "expected_labels": expected_labels,
+                "tags": list(tags or set()),
+                "save_telemetry": telemetry_on,
+                "save_telemetry_synchronously": telemetry_sync,
+                "filters": cast(list[FilterItem], parsed_filters),
+                "prompt": prompt,
+                "use_lookup_cache": use_lookup_cache,
+                "ignore_unlabeled": ignore_unlabeled,
+            },
+            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.")
+
+        predictions = [
+            ClassificationPrediction(
+                prediction_id=prediction["prediction_id"],
+                label=prediction["label"],
+                label_name=prediction["label_name"],
+                score=None,
+                confidence=prediction["confidence"],
+                anomaly_score=prediction["anomaly_score"],
+                memoryset=self.memoryset,
+                model=self,
+                logits=prediction["logits"],
+                input_value=input_value,
+            )
+            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]
+
+    @overload
+    async def apredict(
+        self,
+        value: list[str],
+        expected_labels: list[int] | None = None,
+        filters: list[FilterItemTuple] = [],
+        tags: set[str] | None = None,
+        save_telemetry: TelemetryMode = "on",
+        prompt: str | None = None,
+        use_lookup_cache: bool = True,
+        timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+    ) -> list[ClassificationPrediction]:
+        pass
+
+    @overload
+    async def apredict(
+        self,
+        value: str,
+        expected_labels: int | None = None,
+        filters: list[FilterItemTuple] = [],
+        tags: set[str] | None = None,
+        save_telemetry: TelemetryMode = "on",
+        prompt: str | None = None,
+        use_lookup_cache: bool = True,
+        timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+    ) -> ClassificationPrediction:
+        pass
+
+    async def apredict(
+        self,
+        value: list[str] | str,
+        expected_labels: list[int] | list[str] | int | str | None = None,
+        filters: list[FilterItemTuple] = [],
+        tags: set[str] | None = None,
+        save_telemetry: TelemetryMode = "on",
+        prompt: str | None = None,
+        use_lookup_cache: bool = True,
+        timeout_seconds: int = 10,
+        ignore_unlabeled: bool = False,
+    ) -> list[ClassificationPrediction] | ClassificationPrediction:
+        """
+        Asynchronously predict label(s) for the given input value(s) grounded in similar memories
+
+        Params:
+            value: Value(s) to get predict the labels of
+            expected_labels: Expected label(s) for the given input to record for model evaluation
+            filters: Optional filters to apply during memory lookup
+            tags: Tags to add to the prediction(s)
+            save_telemetry: Whether to save telemetry for the prediction(s). One of
+                * `"off"`: Do not save telemetry
+                * `"on"`: Save telemetry asynchronously unless the `ORCA_SAVE_TELEMETRY_SYNCHRONOUSLY`
+                  environment variable is set.
+                * `"sync"`: Save telemetry synchronously
+                * `"async"`: Save telemetry asynchronously
+            prompt: Optional prompt to use for instruction-tuned embedding models
+            use_lookup_cache: Whether to use cached lookup results for faster predictions
+            timeout_seconds: Timeout in seconds for the request, defaults to 10 seconds
+            ignore_unlabeled: If True, only use labeled memories during lookup.
+                If False (default), allow unlabeled memories when necessary.
+
+        Returns:
+            Label prediction or list of label predictions.
+
+        Raises:
+            ValueError: If timeout_seconds is not a positive integer
+            TimeoutError: If the request times out after the specified duration
+
+        Examples:
+            Predict the label for a single value:
+            >>> prediction = await model.apredict("I am happy", tags={"test"})
+            ClassificationPrediction({label: <positive: 1>, confidence: 0.95, anomaly_score: 0.1, input_value: 'I am happy' })
+
+            Predict the labels for a list of values:
+            >>> predictions = await model.apredict(["I am happy", "I am sad"], expected_labels=[1, 0])
+            [
+                ClassificationPrediction({label: <positive: 1>, confidence: 0.95, anomaly_score: 0.1, input_value: 'I am happy'}),
+                ClassificationPrediction({label: <negative: 0>, confidence: 0.05, anomaly_score: 0.1, input_value: 'I am sad'}),
+            ]
+
+            Using a prompt with an instruction-tuned embedding model:
+            >>> prediction = await model.apredict("I am happy", prompt="Represent this text for sentiment classification:")
+            ClassificationPrediction({label: <positive: 1>, confidence: 0.95, anomaly_score: 0.1, input_value: 'I am happy' })
+        """
+
+        if timeout_seconds <= 0:
+            raise ValueError("timeout_seconds must be a positive integer")
+
+        parsed_filters = [
+            _parse_filter_item_from_tuple(filter) if isinstance(filter, tuple) else filter for filter in filters
+        ]
+
         if any(_is_metric_column(filter[0]) for filter in filters):
             raise ValueError(f"Cannot filter on {filters} - telemetry filters are not supported for predictions")
 
@@ -435,7 +610,8 @@ class ClassificationModel:
             ]
 
         telemetry_on, telemetry_sync = _get_telemetry_config(save_telemetry)
-        response = orca_api.POST(
+        client = OrcaAsyncClient._resolve_client()
+        response = await client.POST(
             "/gpu/classification_model/{name_or_id}/prediction",
             params={"name_or_id": self.id},
             json={
@@ -448,6 +624,7 @@ class ClassificationModel:
                 "filters": cast(list[FilterItem], parsed_filters),
                 "prompt": prompt,
                 "use_lookup_cache": use_lookup_cache,
+                "ignore_unlabeled": ignore_unlabeled,
             },
             timeout=timeout_seconds,
         )
@@ -515,7 +692,8 @@ class ClassificationModel:
             >>> predictions = model.predictions(expected_label_match=False)
             [ClassificationPrediction({label: <positive: 1>, confidence: 0.95, anomaly_score: 0.1, input_value: 'I am happy', expected_label: 0})]
         """
-        predictions = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        predictions = client.POST(
             "/telemetry/prediction",
             json={
                 "model_id": self.id,
@@ -549,9 +727,12 @@ class ClassificationModel:
         label_column: str,
         record_predictions: bool,
         tags: set[str] | None,
+        subsample: int | float | None,
         background: bool = False,
+        ignore_unlabeled: bool = False,
     ) -> ClassificationMetrics | Job[ClassificationMetrics]:
-        response = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        response = client.POST(
             "/classification_model/{model_name_or_id}/evaluation",
             params={"model_name_or_id": self.id},
             json={
@@ -561,13 +742,16 @@ class ClassificationModel:
                 "memoryset_override_name_or_id": self._memoryset_override_id,
                 "record_telemetry": record_predictions,
                 "telemetry_tags": list(tags) if tags else None,
+                "subsample": subsample,
+                "ignore_unlabeled": ignore_unlabeled,
             },
         )
 
         def get_value():
-            res = orca_api.GET(
-                "/classification_model/{model_name_or_id}/evaluation/{task_id}",
-                params={"model_name_or_id": self.id, "task_id": response["task_id"]},
+            client = OrcaClient._resolve_client()
+            res = client.GET(
+                "/classification_model/{model_name_or_id}/evaluation/{job_id}",
+                params={"model_name_or_id": self.id, "job_id": response["job_id"]},
             )
             assert res["result"] is not None
             return ClassificationMetrics(
@@ -584,7 +768,7 @@ class ClassificationModel:
                 roc_curve=res["result"].get("roc_curve"),
             )
 
-        job = Job(response["task_id"], get_value)
+        job = Job(response["job_id"], get_value)
         return job if background else job.result()
 
     def _evaluate_dataset(
@@ -595,6 +779,7 @@ class ClassificationModel:
         record_predictions: bool,
         tags: set[str],
         batch_size: int,
+        ignore_unlabeled: bool,
     ) -> ClassificationMetrics:
         if len(dataset) == 0:
             raise ValueError("Evaluation dataset cannot be empty")
@@ -610,6 +795,7 @@ class ClassificationModel:
                 expected_labels=dataset[i : i + batch_size][label_column],
                 tags=tags,
                 save_telemetry="sync" if record_predictions else "off",
+                ignore_unlabeled=ignore_unlabeled,
             )
         ]
 
@@ -630,7 +816,9 @@ class ClassificationModel:
         record_predictions: bool = False,
         tags: set[str] = {"evaluation"},
         batch_size: int = 100,
+        subsample: int | float | None = None,
         background: Literal[True],
+        ignore_unlabeled: bool = False,
     ) -> Job[ClassificationMetrics]:
         pass
 
@@ -644,7 +832,9 @@ class ClassificationModel:
         record_predictions: bool = False,
         tags: set[str] = {"evaluation"},
         batch_size: int = 100,
+        subsample: int | float | None = None,
         background: Literal[False] = False,
+        ignore_unlabeled: bool = False,
     ) -> ClassificationMetrics:
         pass
 
@@ -657,7 +847,9 @@ class ClassificationModel:
         record_predictions: bool = False,
         tags: set[str] = {"evaluation"},
         batch_size: int = 100,
+        subsample: int | float | None = None,
         background: bool = False,
+        ignore_unlabeled: bool = False,
     ) -> ClassificationMetrics | Job[ClassificationMetrics]:
         """
         Evaluate the classification model on a given dataset or datasource
@@ -669,7 +861,9 @@ class ClassificationModel:
             record_predictions: Whether to record [`ClassificationPrediction`][orca_sdk.telemetry.ClassificationPrediction]s for analysis
             tags: Optional tags to add to the recorded [`ClassificationPrediction`][orca_sdk.telemetry.ClassificationPrediction]s
             batch_size: Batch size for processing Dataset inputs (only used when input is a Dataset)
+            subsample: Optional number (int) of rows to sample or fraction (float in (0, 1]) of data to sample for evaluation.
             background: Whether to run the operation in the background and return a job handle
+            ignore_unlabeled: If True, only use labeled memories during lookup. If False (default), allow unlabeled memories
 
         Returns:
             EvaluationResult containing metrics including accuracy, F1 score, ROC AUC, PR AUC, and anomaly score statistics
@@ -691,7 +885,9 @@ class ClassificationModel:
                 label_column=label_column,
                 record_predictions=record_predictions,
                 tags=tags,
+                subsample=subsample,
                 background=background,
+                ignore_unlabeled=ignore_unlabeled,
             )
         elif isinstance(data, Dataset):
             return self._evaluate_dataset(
@@ -701,6 +897,7 @@ class ClassificationModel:
                 record_predictions=record_predictions,
                 tags=tags,
                 batch_size=batch_size,
+                ignore_unlabeled=ignore_unlabeled,
             )
         else:
             raise ValueError(f"Invalid data type: {type(data)}")
@@ -773,7 +970,8 @@ class ClassificationModel:
             ValueError: If the value does not match previous value types for the category, or is a
                 [`float`][float] that is not between `-1.0` and `+1.0`.
         """
-        orca_api.PUT(
+        client = OrcaClient._resolve_client()
+        client.PUT(
             "/telemetry/prediction/feedback",
             json=[
                 _parse_feedback(f) for f in (cast(list[dict], [feedback]) if isinstance(feedback, dict) else feedback)
@@ -788,7 +986,8 @@ class ClassificationModel:
         num_examples_per_label: int,
         background: bool = False,
     ) -> Job[BootstrappedClassificationModel] | BootstrappedClassificationModel:
-        response = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        response = client.POST(
             "/agents/bootstrap_classification_model",
             json={
                 "model_description": model_description,
@@ -799,11 +998,10 @@ class ClassificationModel:
         )
 
         def get_result() -> BootstrappedClassificationModel:
-            res = orca_api.GET(
-                "/agents/bootstrap_classification_model/{task_id}", params={"task_id": response["task_id"]}
-            )
+            client = OrcaClient._resolve_client()
+            res = client.GET("/agents/bootstrap_classification_model/{job_id}", params={"job_id": response["job_id"]})
             assert res["result"] is not None
             return BootstrappedClassificationModel(res["result"])
 
-        job = Job(response["task_id"], get_result)
+        job = Job(response["job_id"], get_result)
         return job if background else job.result()

orca_sdk/classification_model_test.py

@@ -53,9 +53,10 @@ def test_create_model_already_exists_return(readonly_memoryset, classification_m
     assert new_model.memory_lookup_count == 3
 
 
-def test_create_model_unauthenticated(unauthenticated, readonly_memoryset: LabeledMemoryset):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        ClassificationModel.create("test_model", readonly_memoryset)
+def test_create_model_unauthenticated(unauthenticated_client, readonly_memoryset: LabeledMemoryset):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            ClassificationModel.create("test_model", readonly_memoryset)
 
 
 def test_get_model(classification_model: ClassificationModel):
@@ -68,9 +69,10 @@ def test_get_model(classification_model: ClassificationModel):
     assert fetched_model == classification_model
 
 
-def test_get_model_unauthenticated(unauthenticated):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        ClassificationModel.open("test_model")
+def test_get_model_unauthenticated(unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            ClassificationModel.open("test_model")
 
 
 def test_get_model_invalid_input():
@@ -83,9 +85,10 @@ def test_get_model_not_found():
         ClassificationModel.open(str(uuid4()))
 
 
-def test_get_model_unauthorized(unauthorized, classification_model: ClassificationModel):
-    with pytest.raises(LookupError):
-        ClassificationModel.open(classification_model.name)
+def test_get_model_unauthorized(unauthorized_client, classification_model: ClassificationModel):
+    with unauthorized_client.use():
+        with pytest.raises(LookupError):
+            ClassificationModel.open(classification_model.name)
 
 
 def test_list_models(classification_model: ClassificationModel):
@@ -94,13 +97,15 @@ def test_list_models(classification_model: ClassificationModel):
     assert any(model.name == model.name for model in models)
 
 
-def test_list_models_unauthenticated(unauthenticated):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        ClassificationModel.all()
+def test_list_models_unauthenticated(unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            ClassificationModel.all()
 
 
-def test_list_models_unauthorized(unauthorized, classification_model: ClassificationModel):
-    assert ClassificationModel.all() == []
+def test_list_models_unauthorized(unauthorized_client, classification_model: ClassificationModel):
+    with unauthorized_client.use():
+        assert ClassificationModel.all() == []
 
 
 def test_update_model_attributes(classification_model: ClassificationModel):
@@ -131,9 +136,10 @@ def test_delete_model(readonly_memoryset: LabeledMemoryset):
         ClassificationModel.open("model_to_delete")
 
 
-def test_delete_model_unauthenticated(unauthenticated, classification_model: ClassificationModel):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        ClassificationModel.drop(classification_model.name)
+def test_delete_model_unauthenticated(unauthenticated_client, classification_model: ClassificationModel):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            ClassificationModel.drop(classification_model.name)
 
 
 def test_delete_model_not_found():
@@ -143,9 +149,10 @@ def test_delete_model_not_found():
     ClassificationModel.drop(str(uuid4()), if_not_exists="ignore")
 
 
-def test_delete_model_unauthorized(unauthorized, classification_model: ClassificationModel):
-    with pytest.raises(LookupError):
-        ClassificationModel.drop(classification_model.name)
+def test_delete_model_unauthorized(unauthorized_client, classification_model: ClassificationModel):
+    with unauthorized_client.use():
+        with pytest.raises(LookupError):
+            ClassificationModel.drop(classification_model.name)
 
 
 def test_delete_memoryset_before_model_constraint_violation(hf_dataset):
@@ -254,14 +261,16 @@ def test_predict_disable_telemetry(classification_model: ClassificationModel, la
     assert 0 <= predictions[1].confidence <= 1
 
 
-def test_predict_unauthenticated(unauthenticated, classification_model: ClassificationModel):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        classification_model.predict(["Do you love soup?", "Are cats cute?"])
+def test_predict_unauthenticated(unauthenticated_client, classification_model: ClassificationModel):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            classification_model.predict(["Do you love soup?", "Are cats cute?"])
 
 
-def test_predict_unauthorized(unauthorized, classification_model: ClassificationModel):
-    with pytest.raises(LookupError):
-        classification_model.predict(["Do you love soup?", "Are cats cute?"])
+def test_predict_unauthorized(unauthorized_client, classification_model: ClassificationModel):
+    with unauthorized_client.use():
+        with pytest.raises(LookupError):
+            classification_model.predict(["Do you love soup?", "Are cats cute?"])
 
 
 def test_predict_constraint_violation(readonly_memoryset: LabeledMemoryset):
@@ -396,7 +405,7 @@ def test_last_prediction_with_single(classification_model: ClassificationModel):
 def test_explain(writable_memoryset: LabeledMemoryset):
 
     writable_memoryset.analyze(
-        {"name": "neighbor", "neighbor_counts": [1, 3]},
+        {"name": "distribution", "neighbor_counts": [1, 3]},
         lookup_count=3,
     )
 
@@ -430,7 +439,7 @@ def test_action_recommendation(writable_memoryset: LabeledMemoryset):
     """Test getting action recommendations for predictions"""
 
     writable_memoryset.analyze(
-        {"name": "neighbor", "neighbor_counts": [1, 3]},
+        {"name": "distribution", "neighbor_counts": [1, 3]},
         lookup_count=3,
     )
 
@@ -494,3 +503,62 @@ def test_predict_with_prompt(classification_model: ClassificationModel):
     # Both should work and return valid predictions
     assert prediction_with_prompt.label is not None
     assert prediction_without_prompt.label is not None
+
+
+@pytest.mark.asyncio
+async def test_predict_async_single(classification_model: ClassificationModel, label_names: list[str]):
+    """Test async prediction with a single value"""
+    prediction = await classification_model.apredict("Do you love soup?")
+    assert isinstance(prediction, ClassificationPrediction)
+    assert prediction.prediction_id is not None
+    assert prediction.label == 0
+    assert prediction.label_name == label_names[0]
+    assert 0 <= prediction.confidence <= 1
+    assert prediction.logits is not None
+    assert len(prediction.logits) == 2
+
+
+@pytest.mark.asyncio
+async def test_predict_async_batch(classification_model: ClassificationModel, label_names: list[str]):
+    """Test async prediction with a batch of values"""
+    predictions = await classification_model.apredict(["Do you love soup?", "Are cats cute?"])
+    assert len(predictions) == 2
+    assert predictions[0].prediction_id is not None
+    assert predictions[1].prediction_id is not None
+    assert predictions[0].label == 0
+    assert predictions[0].label_name == label_names[0]
+    assert 0 <= predictions[0].confidence <= 1
+    assert predictions[1].label == 1
+    assert predictions[1].label_name == label_names[1]
+    assert 0 <= predictions[1].confidence <= 1
+
+
+@pytest.mark.asyncio
+async def test_predict_async_with_expected_labels(classification_model: ClassificationModel):
+    """Test async prediction with expected labels"""
+    prediction = await classification_model.apredict("Do you love soup?", expected_labels=1)
+    assert prediction.expected_label == 1
+
+
+@pytest.mark.asyncio
+async def test_predict_async_disable_telemetry(classification_model: ClassificationModel, label_names: list[str]):
+    """Test async prediction with telemetry disabled"""
+    predictions = await classification_model.apredict(["Do you love soup?", "Are cats cute?"], save_telemetry="off")
+    assert len(predictions) == 2
+    assert predictions[0].prediction_id is None
+    assert predictions[1].prediction_id is None
+    assert predictions[0].label == 0
+    assert predictions[0].label_name == label_names[0]
+    assert 0 <= predictions[0].confidence <= 1
+    assert predictions[1].label == 1
+    assert predictions[1].label_name == label_names[1]
+    assert 0 <= predictions[1].confidence <= 1
+
+
+@pytest.mark.asyncio
+async def test_predict_async_with_filters(classification_model: ClassificationModel):
+    """Test async prediction with filters"""
+    # there are no memories with label 0 and key g2, so we force a wrong prediction
+    filtered_prediction = await classification_model.apredict("I love soup", filters=[("key", "==", "g2")])
+    assert filtered_prediction.label == 1
+    assert filtered_prediction.label_name == "cats"