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

orca_sdk/classification_model.py

@@ -3,19 +3,27 @@ from __future__ import annotations
 import logging
 from contextlib import contextmanager
 from datetime import datetime
-from typing import Any, Generator, Iterable, Literal, cast, overload
+from typing import (
+    Any,
+    Generator,
+    Iterable,
+    Literal,
+    cast,
+    overload,
+)
 
 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,
     ClassificationModelMetadata,
+    OrcaClient,
     PredictiveModelUpdate,
     RACHeadType,
-    orca_api,
 )
 from .datasource import Datasource
 from .job import Job
@@ -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:
@@ -435,7 +448,150 @@ class ClassificationModel:
             ]
 
         telemetry_on, telemetry_sync = _get_telemetry_config(save_telemetry)
-        response = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        response = client.POST(
+            "/gpu/classification_model/{name_or_id}/prediction",
+            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,
+            },
+            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,
+    ) -> 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,
+    ) -> 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,
+    ) -> 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
+
+        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")
+
+        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
+            ]
+
+        telemetry_on, telemetry_sync = _get_telemetry_config(save_telemetry)
+        client = OrcaAsyncClient._resolve_client()
+        response = await client.POST(
             "/gpu/classification_model/{name_or_id}/prediction",
             params={"name_or_id": self.id},
             json={
@@ -515,7 +671,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,
@@ -551,7 +708,8 @@ class ClassificationModel:
         tags: set[str] | None,
         background: 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={
@@ -565,7 +723,8 @@ class ClassificationModel:
         )
 
         def get_value():
-            res = orca_api.GET(
+            client = OrcaClient._resolve_client()
+            res = client.GET(
                 "/classification_model/{model_name_or_id}/evaluation/{task_id}",
                 params={"model_name_or_id": self.id, "task_id": response["task_id"]},
             )
@@ -773,7 +932,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 +948,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,7 +960,8 @@ class ClassificationModel:
         )
 
         def get_result() -> BootstrappedClassificationModel:
-            res = orca_api.GET(
+            client = OrcaClient._resolve_client()
+            res = client.GET(
                 "/agents/bootstrap_classification_model/{task_id}", params={"task_id": response["task_id"]}
             )
             assert res["result"] is not None

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"