orca-sdk 0.0.96__py3-none-any.whl → 0.0.98__py3-none-any.whl

orca_sdk/classification_model.py

@@ -1,56 +1,61 @@
 from __future__ import annotations
 
 import logging
-import os
 from contextlib import contextmanager
 from datetime import datetime
 from typing import Any, Generator, Iterable, Literal, cast, overload
-from uuid import UUID
 
 from datasets import Dataset
 
-from ._generated_api_client.api import (
-    create_classification_model_gpu,
-    delete_classification_model,
-    evaluate_classification_model,
-    get_classification_model,
-    get_classification_model_evaluation,
-    list_classification_models,
-    list_predictions,
-    predict_label_gpu,
-    record_prediction_feedback,
-    update_classification_model,
-)
-from ._generated_api_client.models import (
-    ClassificationEvaluationRequest,
+from ._shared.metrics import ClassificationMetrics, calculate_classification_metrics
+from ._utils.common import UNSET, CreateMode, DropMode
+from .client import (
+    BootstrapClassificationModelMeta,
+    BootstrapClassificationModelResult,
     ClassificationModelMetadata,
-    ClassificationPredictionRequest,
-    CreateClassificationModelRequest,
-    LabelPredictionWithMemoriesAndFeedback,
-    ListPredictionsRequest,
-)
-from ._generated_api_client.models import (
-    PredictionSortItemItemType0 as PredictionSortColumns,
-)
-from ._generated_api_client.models import (
-    PredictionSortItemItemType1 as PredictionSortDirection,
-)
-from ._generated_api_client.models import (
     PredictiveModelUpdate,
     RACHeadType,
+    orca_api,
 )
-from ._generated_api_client.types import UNSET as CLIENT_UNSET
-from ._shared.metrics import ClassificationMetrics, calculate_classification_metrics
-from ._utils.common import UNSET, CreateMode, DropMode
 from .datasource import Datasource
 from .job import Job
 from .memoryset import (
     FilterItem,
     FilterItemTuple,
     LabeledMemoryset,
+    _is_metric_column,
     _parse_filter_item_from_tuple,
 )
-from .telemetry import ClassificationPrediction, _parse_feedback
+from .telemetry import (
+    ClassificationPrediction,
+    TelemetryMode,
+    _get_telemetry_config,
+    _parse_feedback,
+)
+
+
+class BootstrappedClassificationModel:
+
+    datasource: Datasource | None
+    memoryset: LabeledMemoryset | None
+    classification_model: ClassificationModel | None
+    agent_output: BootstrapClassificationModelResult | None
+
+    def __init__(self, metadata: BootstrapClassificationModelMeta):
+        self.datasource = Datasource.open(metadata["datasource_meta"]["id"])
+        self.memoryset = LabeledMemoryset.open(metadata["memoryset_meta"]["id"])
+        self.classification_model = ClassificationModel.open(metadata["model_meta"]["id"])
+        self.agent_output = metadata["agent_output"]
+
+    def __repr__(self):
+        return (
+            "BootstrappedClassificationModel({\n"
+            f"    datasource: {self.datasource},\n"
+            f"    memoryset: {self.memoryset},\n"
+            f"    classification_model: {self.classification_model},\n"
+            f"    agent_output: {self.agent_output},\n"
+            "})"
+        )
 
 
 class ClassificationModel:
@@ -86,18 +91,18 @@ class ClassificationModel:
 
     def __init__(self, metadata: ClassificationModelMetadata):
         # for internal use only, do not document
-        self.id = metadata.id
-        self.name = metadata.name
-        self.description = metadata.description
-        self.memoryset = LabeledMemoryset.open(metadata.memoryset_id)
-        self.head_type = metadata.head_type
-        self.num_classes = metadata.num_classes
-        self.memory_lookup_count = metadata.memory_lookup_count
-        self.weigh_memories = metadata.weigh_memories
-        self.min_memory_weight = metadata.min_memory_weight
-        self.version = metadata.version
-        self.locked = metadata.locked
-        self.created_at = metadata.created_at
+        self.id = metadata["id"]
+        self.name = metadata["name"]
+        self.description = metadata["description"]
+        self.memoryset = LabeledMemoryset.open(metadata["memoryset_id"])
+        self.head_type = metadata["head_type"]
+        self.num_classes = metadata["num_classes"]
+        self.memory_lookup_count = metadata["memory_lookup_count"]
+        self.weigh_memories = metadata["weigh_memories"]
+        self.min_memory_weight = metadata["min_memory_weight"]
+        self.version = metadata["version"]
+        self.locked = metadata["locked"]
+        self.created_at = datetime.fromisoformat(metadata["created_at"])
 
         self._memoryset_override_id: str | None = None
         self._last_prediction: ClassificationPrediction | None = None
@@ -140,7 +145,7 @@ class ClassificationModel:
         cls,
         name: str,
         memoryset: LabeledMemoryset,
-        head_type: Literal["BMMOE", "FF", "KNN", "MMOE"] = "KNN",
+        head_type: RACHeadType = "KNN",
         *,
         description: str | None = None,
         num_classes: int | None = None,
@@ -206,17 +211,18 @@ class ClassificationModel:
 
                 return existing
 
-        metadata = create_classification_model_gpu(
-            body=CreateClassificationModelRequest(
-                name=name,
-                memoryset_id=memoryset.id,
-                head_type=RACHeadType(head_type),
-                memory_lookup_count=memory_lookup_count,
-                num_classes=num_classes,
-                weigh_memories=weigh_memories,
-                min_memory_weight=min_memory_weight,
-                description=description,
-            ),
+        metadata = orca_api.POST(
+            "/classification_model",
+            json={
+                "name": name,
+                "memoryset_name_or_id": memoryset.id,
+                "head_type": head_type,
+                "memory_lookup_count": memory_lookup_count,
+                "num_classes": num_classes,
+                "weigh_memories": weigh_memories,
+                "min_memory_weight": min_memory_weight,
+                "description": description,
+            },
         )
         return cls(metadata)
 
@@ -234,7 +240,7 @@ class ClassificationModel:
         Raises:
             LookupError: If the classification model does not exist
         """
-        return cls(get_classification_model(name))
+        return cls(orca_api.GET("/classification_model/{name_or_id}", params={"name_or_id": name}))
 
     @classmethod
     def exists(cls, name_or_id: str) -> bool:
@@ -261,7 +267,7 @@ class ClassificationModel:
         Returns:
             List of handles to all classification models in the OrcaCloud
         """
-        return [cls(metadata) for metadata in list_classification_models()]
+        return [cls(metadata) for metadata in orca_api.GET("/classification_model")]
 
     @classmethod
     def drop(cls, name_or_id: str, if_not_exists: DropMode = "error"):
@@ -280,7 +286,7 @@ class ClassificationModel:
             LookupError: If the classification model does not exist and if_not_exists is `"error"`
         """
         try:
-            delete_classification_model(name_or_id)
+            orca_api.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":
@@ -311,11 +317,12 @@ class ClassificationModel:
             Lock the model:
             >>> model.set(locked=True)
         """
-        update_data = PredictiveModelUpdate(
-            description=CLIENT_UNSET if description is UNSET else description,
-            locked=CLIENT_UNSET if locked is UNSET else locked,
-        )
-        update_classification_model(self.id, body=update_data)
+        update: PredictiveModelUpdate = {}
+        if description is not UNSET:
+            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)
         self.refresh()
 
     def lock(self) -> None:
@@ -333,7 +340,9 @@ class ClassificationModel:
         expected_labels: list[int] | None = None,
         filters: list[FilterItemTuple] = [],
         tags: set[str] | None = None,
-        save_telemetry: Literal["off", "on", "sync", "async"] = "on",
+        save_telemetry: TelemetryMode = "on",
+        prompt: str | None = None,
+        use_lookup_cache: bool = True,
     ) -> list[ClassificationPrediction]:
         pass
 
@@ -344,17 +353,21 @@ class ClassificationModel:
         expected_labels: int | None = None,
         filters: list[FilterItemTuple] = [],
         tags: set[str] | None = None,
-        save_telemetry: Literal["off", "on", "sync", "async"] = "on",
+        save_telemetry: TelemetryMode = "on",
+        prompt: str | None = None,
+        use_lookup_cache: bool = True,
     ) -> ClassificationPrediction:
         pass
 
     def predict(
         self,
         value: list[str] | str,
-        expected_labels: list[int] | int | None = None,
+        expected_labels: list[int] | list[str] | int | str | None = None,
         filters: list[FilterItemTuple] = [],
         tags: set[str] | None = None,
-        save_telemetry: Literal["off", "on", "sync", "async"] = "on",
+        save_telemetry: TelemetryMode = "on",
+        prompt: str | None = None,
+        use_lookup_cache: bool = True,
     ) -> list[ClassificationPrediction] | ClassificationPrediction:
         """
         Predict label(s) for the given input value(s) grounded in similar memories
@@ -362,6 +375,7 @@ class ClassificationModel:
         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
@@ -369,6 +383,7 @@ class ClassificationModel:
                   environment variable is set.
                 * `"sync"`: Save telemetry synchronously
                 * `"async"`: Save telemetry asynchronously
+            prompt: Optional prompt to use for instruction-tuned embedding models
 
         Returns:
             Label prediction or list of label predictions
@@ -384,48 +399,60 @@ class ClassificationModel:
                 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 = model.predict("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' })
         """
 
         parsed_filters = [
             _parse_filter_item_from_tuple(filter) if isinstance(filter, tuple) else filter for filter in filters
         ]
 
-        if not all(isinstance(filter, FilterItem) for filter in parsed_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")
 
-        response = predict_label_gpu(
-            self.id,
-            body=ClassificationPredictionRequest(
-                input_values=value if isinstance(value, list) else [value],
-                memoryset_override_id=self._memoryset_override_id,
-                expected_labels=(
-                    expected_labels
-                    if isinstance(expected_labels, list)
-                    else [expected_labels] if expected_labels is not None else None
-                ),
-                tags=list(tags or set()),
-                save_telemetry=save_telemetry != "off",
-                save_telemetry_synchronously=(
-                    os.getenv("ORCA_SAVE_TELEMETRY_SYNCHRONOUSLY", "0") != "0" or save_telemetry == "sync"
-                ),
-                filters=cast(list[FilterItem], parsed_filters),
-            ),
+        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)
+        response = orca_api.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,
+            },
         )
 
-        if save_telemetry != "off" and any(p.prediction_id is None for p in response):
+        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,
+                prediction_id=prediction["prediction_id"],
+                label=prediction["label"],
+                label_name=prediction["label_name"],
                 score=None,
-                confidence=prediction.confidence,
-                anomaly_score=prediction.anomaly_score,
+                confidence=prediction["confidence"],
+                anomaly_score=prediction["anomaly_score"],
                 memoryset=self.memoryset,
                 model=self,
-                logits=prediction.logits,
+                logits=prediction["logits"],
                 input_value=input_value,
             )
             for prediction, input_value in zip(response, value if isinstance(value, list) else [value])
@@ -439,7 +466,7 @@ class ClassificationModel:
         limit: int = 100,
         offset: int = 0,
         tag: str | None = None,
-        sort: list[tuple[PredictionSortColumns, PredictionSortDirection]] = [],
+        sort: list[tuple[Literal["anomaly_score", "confidence", "timestamp"], Literal["asc", "desc"]]] = [],
         expected_label_match: bool | None = None,
     ) -> list[ClassificationPrediction]:
         """
@@ -475,30 +502,31 @@ 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 = list_predictions(
-            body=ListPredictionsRequest(
-                model_id=self.id,
-                limit=limit,
-                offset=offset,
-                sort=cast(list[list[PredictionSortColumns | PredictionSortDirection]], sort),
-                tag=tag,
-                expected_label_match=expected_label_match,
-            ),
+        predictions = orca_api.POST(
+            "/telemetry/prediction",
+            json={
+                "model_id": self.id,
+                "limit": limit,
+                "offset": offset,
+                "sort": [list(sort_item) for sort_item in sort],
+                "tag": tag,
+                "expected_label_match": expected_label_match,
+            },
         )
         return [
             ClassificationPrediction(
-                prediction_id=prediction.prediction_id,
-                label=prediction.label,
-                label_name=prediction.label_name,
+                prediction_id=prediction["prediction_id"],
+                label=prediction["label"],
+                label_name=prediction["label_name"],
                 score=None,
-                confidence=prediction.confidence,
-                anomaly_score=prediction.anomaly_score,
+                confidence=prediction["confidence"],
+                anomaly_score=prediction["anomaly_score"],
                 memoryset=self.memoryset,
                 model=self,
                 telemetry=prediction,
             )
             for prediction in predictions
-            if isinstance(prediction, LabelPredictionWithMemoriesAndFeedback)
+            if "label" in prediction
         ]
 
     def _evaluate_datasource(
@@ -510,23 +538,28 @@ class ClassificationModel:
         tags: set[str] | None,
         background: bool = False,
     ) -> ClassificationMetrics | Job[ClassificationMetrics]:
-        response = evaluate_classification_model(
-            self.id,
-            body=ClassificationEvaluationRequest(
-                datasource_id=datasource.id,
-                datasource_label_column=label_column,
-                datasource_value_column=value_column,
-                memoryset_override_id=self._memoryset_override_id,
-                record_telemetry=record_predictions,
-                telemetry_tags=list(tags) if tags else None,
-            ),
+        response = orca_api.POST(
+            "/classification_model/{model_name_or_id}/evaluation",
+            params={"model_name_or_id": self.id},
+            json={
+                "datasource_name_or_id": datasource.id,
+                "datasource_label_column": label_column,
+                "datasource_value_column": value_column,
+                "memoryset_override_name_or_id": self._memoryset_override_id,
+                "record_telemetry": record_predictions,
+                "telemetry_tags": list(tags) if tags else None,
+            },
         )
 
-        job = Job(
-            response.task_id,
-            lambda: (r := get_classification_model_evaluation(self.id, UUID(response.task_id)).result)
-            and ClassificationMetrics(**r.to_dict()),
-        )
+        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"]},
+            )
+            assert res["result"] is not None
+            return ClassificationMetrics(**res["result"])
+
+        job = Job(response["task_id"], get_value)
         return job if background else job.result()
 
     def _evaluate_dataset(
@@ -709,8 +742,37 @@ 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`.
         """
-        record_prediction_feedback(
-            body=[
+        orca_api.PUT(
+            "/telemetry/prediction/feedback",
+            json=[
                 _parse_feedback(f) for f in (cast(list[dict], [feedback]) if isinstance(feedback, dict) else feedback)
             ],
         )
+
+    @staticmethod
+    def bootstrap_model(
+        model_description: str,
+        label_names: list[str],
+        initial_examples: list[tuple[str, str]],
+        num_examples_per_label: int,
+        background: bool = False,
+    ) -> Job[BootstrappedClassificationModel] | BootstrappedClassificationModel:
+        response = orca_api.POST(
+            "/agents/bootstrap_classification_model",
+            json={
+                "model_description": model_description,
+                "label_names": label_names,
+                "initial_examples": [{"text": text, "label_name": label_name} for text, label_name in initial_examples],
+                "num_examples_per_label": num_examples_per_label,
+            },
+        )
+
+        def get_result() -> BootstrappedClassificationModel:
+            res = orca_api.GET(
+                "/agents/bootstrap_classification_model/{task_id}", params={"task_id": response["task_id"]}
+            )
+            assert res["result"] is not None
+            return BootstrappedClassificationModel(res["result"])
+
+        job = Job(response["task_id"], get_result)
+        return job if background else job.result()

orca_sdk/classification_model_test.py

@@ -326,6 +326,37 @@ def test_predict_with_filters(classification_model: ClassificationModel):
     assert filtered_prediction.label_name == "cats"
 
 
+def test_predict_with_memoryset_update(writable_memoryset: LabeledMemoryset):
+    model = ClassificationModel.create(
+        "test_predict_with_memoryset_update",
+        writable_memoryset,
+        num_classes=2,
+        memory_lookup_count=3,
+    )
+
+    prediction = model.predict("Do you love soup?")
+    assert prediction.label == 0
+    assert prediction.label_name == "soup"
+
+    # insert new memories
+    writable_memoryset.insert(
+        [
+            {"value": "Do you love soup?", "label": 1, "key": "g1"},
+            {"value": "Do you love soup for dinner?", "label": 1, "key": "g2"},
+            {"value": "Do you love crackers?", "label": 1, "key": "g2"},
+            {"value": "Do you love broth?", "label": 1, "key": "g2"},
+            {"value": "Do you love chicken soup?", "label": 1, "key": "g2"},
+            {"value": "Do you love chicken soup for dinner?", "label": 1, "key": "g2"},
+            {"value": "Do you love chicken soup for dinner?", "label": 1, "key": "g2"},
+        ],
+    )
+    prediction = model.predict("Do you love soup?")
+    assert prediction.label == 1
+    assert prediction.label_name == "cats"
+
+    ClassificationModel.drop("test_predict_with_memoryset_update")
+
+
 def test_last_prediction_with_batch(classification_model: ClassificationModel):
     predictions = classification_model.predict(["Do you love soup?", "Are cats cute?"])
     assert classification_model.last_prediction is not None
@@ -374,3 +405,74 @@ def test_explain(writable_memoryset: LabeledMemoryset):
             raise e
     finally:
         ClassificationModel.drop("test_model_for_explain")
+
+
+@skip_in_ci("We don't have Anthropic API key in CI")
+def test_action_recommendation(writable_memoryset: LabeledMemoryset):
+    """Test getting action recommendations for predictions"""
+
+    writable_memoryset.analyze(
+        {"name": "neighbor", "neighbor_counts": [1, 3]},
+        lookup_count=3,
+    )
+
+    model = ClassificationModel.create(
+        "test_model_for_action",
+        writable_memoryset,
+        num_classes=2,
+        memory_lookup_count=3,
+        description="This is a test model for action recommendations",
+    )
+
+    # Make a prediction with expected label to simulate incorrect prediction
+    prediction = model.predict("Do you love soup?", expected_labels=1)
+
+    memoryset_length = model.memoryset.length
+
+    try:
+        # Get action recommendation
+        action, rationale = prediction.recommend_action()
+
+        assert action is not None
+        assert rationale is not None
+        assert action in ["remove_duplicates", "detect_mislabels", "add_memories", "finetuning"]
+        assert len(rationale) > 10
+
+        # Test memory suggestions
+        suggestions_response = prediction.generate_memory_suggestions(num_memories=2)
+        memory_suggestions = suggestions_response.suggestions
+
+        assert memory_suggestions is not None
+        assert len(memory_suggestions) == 2
+
+        for suggestion in memory_suggestions:
+            assert isinstance(suggestion[0], str)
+            assert len(suggestion[0]) > 0
+            assert isinstance(suggestion[1], str)
+            assert suggestion[1] in model.memoryset.label_names
+
+        suggestions_response.apply()
+
+        model.memoryset.refresh()
+        assert model.memoryset.length == memoryset_length + 2
+
+    except Exception as e:
+        if "ANTHROPIC_API_KEY" in str(e):
+            logging.info("Skipping agent tests because ANTHROPIC_API_KEY is not set")
+        else:
+            raise e
+    finally:
+        ClassificationModel.drop("test_model_for_action")
+
+
+def test_predict_with_prompt(classification_model: ClassificationModel):
+    """Test that prompt parameter is properly passed through to predictions"""
+    # Test with an instruction-supporting embedding model if available
+    prediction_with_prompt = classification_model.predict(
+        "I love this product!", prompt="Represent this text for sentiment classification:"
+    )
+    prediction_without_prompt = classification_model.predict("I love this product!")
+
+    # Both should work and return valid predictions
+    assert prediction_with_prompt.label is not None
+    assert prediction_without_prompt.label is not None