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

orca_sdk/memoryset_test.py

@@ -1,4 +1,3 @@
-import os
 import random
 from uuid import uuid4
 
@@ -364,17 +363,14 @@ def test_clone_memoryset(readonly_memoryset: LabeledMemoryset):
 
 
 def test_embedding_evaluation(eval_datasource: Datasource):
-    response = LabeledMemoryset.run_embedding_evaluation(
+    results = LabeledMemoryset.run_embedding_evaluation(
         eval_datasource, embedding_models=["CDE_SMALL"], neighbor_count=2
     )
-    assert response is not None
-    assert isinstance(response, dict)
-    assert response is not None
-    assert isinstance(response["evaluation_results"], list)
-    assert len(response["evaluation_results"]) == 1
-    assert response["evaluation_results"][0] is not None
-    assert response["evaluation_results"][0]["embedding_model_name"] == "CDE_SMALL"
-    assert response["evaluation_results"][0]["embedding_model_path"] == "OrcaDB/cde-small-v1"
+    assert isinstance(results, list)
+    assert len(results) == 1
+    assert results[0] is not None
+    assert results[0]["embedding_model_name"] == "CDE_SMALL"
+    assert results[0]["embedding_model_path"] == "OrcaDB/cde-small-v1"
 
 
 @pytest.fixture(scope="function")
@@ -427,7 +423,7 @@ def test_get_cascading_edits_suggestions(writable_memoryset: LabeledMemoryset):
 
     # Validate the suggestions
     assert len(suggestions) == 1
-    assert suggestions[0].neighbor.value == mislabeled_soup_text
+    assert suggestions[0]["neighbor"]["value"] == mislabeled_soup_text
 
 
 def test_analyze_invalid_analysis_name(readonly_memoryset: LabeledMemoryset):

orca_sdk/regression_model.py

@@ -1,53 +1,29 @@
 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
 
-import numpy as np
 from datasets import Dataset
 
-from ._generated_api_client.api import (
-    create_regression_model_gpu,
-    delete_regression_model,
-    evaluate_regression_model,
-    get_regression_model,
-    get_regression_model_evaluation,
-    list_predictions,
-    list_regression_models,
-    predict_score_gpu,
-    record_prediction_feedback,
-    update_regression_model,
-)
-from ._generated_api_client.models import (
-    CreateRegressionModelRequest,
-    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 (
+from ._shared.metrics import RegressionMetrics, calculate_regression_metrics
+from ._utils.common import UNSET, CreateMode, DropMode
+from .client import (
     PredictiveModelUpdate,
     RARHeadType,
-    RegressionEvaluationRequest,
     RegressionModelMetadata,
-    RegressionPredictionRequest,
-    ScorePredictionWithMemoriesAndFeedback,
+    orca_api,
 )
-from ._generated_api_client.types import UNSET as CLIENT_UNSET
-from ._generated_api_client.types import Response
-from ._shared.metrics import RegressionMetrics, calculate_regression_metrics
-from ._utils.common import UNSET, CreateMode, DropMode
 from .datasource import Datasource
 from .job import Job
 from .memoryset import ScoredMemoryset
-from .telemetry import FeedbackCategory, RegressionPrediction, _parse_feedback
+from .telemetry import (
+    RegressionPrediction,
+    TelemetryMode,
+    _get_telemetry_config,
+    _parse_feedback,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -86,17 +62,17 @@ class RegressionModel:
 
     def __init__(self, metadata: RegressionModelMetadata):
         # for internal use only, do not document
-        self.id = metadata.id
-        self.name = metadata.name
-        self.description = metadata.description
-        self.memoryset = ScoredMemoryset.open(metadata.memoryset_id)
-        self.head_type = metadata.head_type
-        self.memory_lookup_count = metadata.memory_lookup_count
-        self.version = metadata.version
-        self.locked = metadata.locked
-        self.created_at = metadata.created_at
-        self.updated_at = metadata.updated_at
-        self.memoryset_id = metadata.memoryset_id
+        self.id = metadata["id"]
+        self.name = metadata["name"]
+        self.description = metadata["description"]
+        self.memoryset = ScoredMemoryset.open(metadata["memoryset_id"])
+        self.head_type = metadata["head_type"]
+        self.memory_lookup_count = metadata["memory_lookup_count"]
+        self.version = metadata["version"]
+        self.locked = metadata["locked"]
+        self.created_at = datetime.fromisoformat(metadata["created_at"])
+        self.updated_at = datetime.fromisoformat(metadata["updated_at"])
+        self.memoryset_id = metadata["memoryset_id"]
 
         self._memoryset_override_id = None
         self._last_prediction = None
@@ -178,13 +154,14 @@ class RegressionModel:
 
                 return existing
 
-        metadata = create_regression_model_gpu(
-            body=CreateRegressionModelRequest(
-                name=name,
-                memoryset_id=memoryset.id,
-                memory_lookup_count=memory_lookup_count,
-                description=description,
-            )
+        metadata = orca_api.POST(
+            "/regression_model",
+            json={
+                "name": name,
+                "memoryset_name_or_id": memoryset.id,
+                "memory_lookup_count": memory_lookup_count,
+                "description": description,
+            },
         )
         return cls(metadata)
 
@@ -202,7 +179,7 @@ class RegressionModel:
         Raises:
             LookupError: If the regression model does not exist
         """
-        return cls(get_regression_model(name))
+        return cls(orca_api.GET("/regression_model/{name_or_id}", params={"name_or_id": name}))
 
     @classmethod
     def exists(cls, name_or_id: str) -> bool:
@@ -229,7 +206,7 @@ class RegressionModel:
         Returns:
             List of handles to all regression models in the OrcaCloud
         """
-        return [cls(metadata) for metadata in list_regression_models()]
+        return [cls(metadata) for metadata in orca_api.GET("/regression_model")]
 
     @classmethod
     def drop(cls, name_or_id: str, if_not_exists: DropMode = "error"):
@@ -248,7 +225,7 @@ class RegressionModel:
             LookupError: If the regression model does not exist and if_not_exists is `"error"`
         """
         try:
-            delete_regression_model(name_or_id)
+            orca_api.DELETE("/regression_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":
@@ -279,11 +256,12 @@ class RegressionModel:
             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_regression_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("/regression_model/{name_or_id}", params={"name_or_id": self.id}, json=update)
         self.refresh()
 
     def lock(self) -> None:
@@ -300,7 +278,9 @@ class RegressionModel:
         value: str,
         expected_scores: float | None = None,
         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,
     ) -> RegressionPrediction: ...
 
     @overload
@@ -309,7 +289,9 @@ class RegressionModel:
         value: list[str],
         expected_scores: list[float] | None = None,
         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[RegressionPrediction]: ...
 
     # TODO: add filter support
@@ -318,7 +300,9 @@ class RegressionModel:
         value: str | list[str],
         expected_scores: float | list[float] | None = None,
         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,
     ) -> RegressionPrediction | list[RegressionPrediction]:
         """
         Make predictions using the regression model.
@@ -331,6 +315,7 @@ class RegressionModel:
                 which will save telemetry asynchronously unless the `ORCA_SAVE_TELEMETRY_SYNCHRONOUSLY`
                 environment variable is set to `"1"`. You can also pass `"sync"` or `"async"` to
                 explicitly set the save mode.
+            prompt: Optional prompt for instruction-tuned embedding models
 
         Returns:
             Single RegressionPrediction or list of RegressionPrediction objects
@@ -338,35 +323,37 @@ class RegressionModel:
         Raises:
             ValueError: If expected_scores length doesn't match value length for batch predictions
         """
-        response = predict_score_gpu(
-            name_or_id=self.name,
-            body=RegressionPredictionRequest(
-                input_values=value if isinstance(value, list) else [value],
-                expected_scores=(
+        telemetry_on, telemetry_sync = _get_telemetry_config(save_telemetry)
+        response = orca_api.POST(
+            "/gpu/regression_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_scores": (
                     expected_scores
                     if isinstance(expected_scores, list)
                     else [expected_scores] if expected_scores is not None else None
                 ),
-                memoryset_override_id=self._memoryset_override_id,
-                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"
-                ),
-            ),
+                "tags": list(tags or set()),
+                "save_telemetry": telemetry_on,
+                "save_telemetry_synchronously": telemetry_sync,
+                "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 = [
             RegressionPrediction(
-                prediction_id=prediction.prediction_id,
+                prediction_id=prediction["prediction_id"],
                 label=None,
                 label_name=None,
-                score=prediction.score,
-                confidence=prediction.confidence,
-                anomaly_score=prediction.anomaly_score,
+                score=prediction["score"],
+                confidence=prediction["confidence"],
+                anomaly_score=prediction["anomaly_score"],
                 memoryset=self.memoryset,
                 model=self,
                 logits=None,
@@ -383,7 +370,7 @@ class RegressionModel:
         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"]]] = [],
     ) -> list[RegressionPrediction]:
         """
         Get a list of predictions made by this model
@@ -411,23 +398,24 @@ 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'})]
         """
-        predictions = list_predictions(
-            body=ListPredictionsRequest(
-                model_id=self.id,
-                limit=limit,
-                offset=offset,
-                sort=cast(list[list[PredictionSortColumns | PredictionSortDirection]], sort),
-                tag=tag,
-            ),
+        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,
+            },
         )
         return [
             RegressionPrediction(
-                prediction_id=prediction.prediction_id,
+                prediction_id=prediction["prediction_id"],
                 label=None,
                 label_name=None,
-                score=prediction.score,
-                confidence=prediction.confidence,
-                anomaly_score=prediction.anomaly_score,
+                score=prediction["score"],
+                confidence=prediction["confidence"],
+                anomaly_score=prediction["anomaly_score"],
                 memoryset=self.memoryset,
                 model=self,
                 telemetry=prediction,
@@ -435,7 +423,7 @@ class RegressionModel:
                 input_value=None,
             )
             for prediction in predictions
-            if isinstance(prediction, ScorePredictionWithMemoriesAndFeedback)
+            if "score" in prediction
         ]
 
     def _evaluate_datasource(
@@ -447,23 +435,28 @@ class RegressionModel:
         tags: set[str] | None,
         background: bool = False,
     ) -> RegressionMetrics | Job[RegressionMetrics]:
-        response = evaluate_regression_model(
-            self.id,
-            body=RegressionEvaluationRequest(
-                datasource_id=datasource.id,
-                datasource_score_column=score_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(
+            "/regression_model/{model_name_or_id}/evaluation",
+            params={"model_name_or_id": self.id},
+            json={
+                "datasource_name_or_id": datasource.id,
+                "datasource_score_column": score_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_regression_model_evaluation(self.id, UUID(response.task_id)).result)
-            and RegressionMetrics(**r.to_dict()),
-        )
+        def get_value():
+            res = orca_api.GET(
+                "/regression_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 RegressionMetrics(**res["result"])
+
+        job = Job(response["task_id"], get_value)
         return job if background else job.result()
 
     def _evaluate_dataset(
@@ -474,6 +467,7 @@ class RegressionModel:
         record_predictions: bool,
         tags: set[str],
         batch_size: int,
+        prompt: str | None = None,
     ) -> RegressionMetrics:
         predictions = [
             prediction
@@ -483,6 +477,7 @@ class RegressionModel:
                 expected_scores=dataset[i : i + batch_size][score_column],
                 tags=tags,
                 save_telemetry="sync" if record_predictions else "off",
+                prompt=prompt,
             )
         ]
 
@@ -502,6 +497,7 @@ class RegressionModel:
         record_predictions: bool = False,
         tags: set[str] = {"evaluation"},
         batch_size: int = 100,
+        prompt: str | None = None,
         background: Literal[True],
     ) -> Job[RegressionMetrics]:
         pass
@@ -516,6 +512,7 @@ class RegressionModel:
         record_predictions: bool = False,
         tags: set[str] = {"evaluation"},
         batch_size: int = 100,
+        prompt: str | None = None,
         background: Literal[False] = False,
     ) -> RegressionMetrics:
         pass
@@ -529,6 +526,7 @@ class RegressionModel:
         record_predictions: bool = False,
         tags: set[str] = {"evaluation"},
         batch_size: int = 100,
+        prompt: str | None = None,
         background: bool = False,
     ) -> RegressionMetrics | Job[RegressionMetrics]:
         """
@@ -541,6 +539,7 @@ class RegressionModel:
             record_predictions: Whether to record [`RegressionPrediction`][orca_sdk.telemetry.RegressionPrediction]s for analysis
             tags: Optional tags to add to the recorded [`RegressionPrediction`][orca_sdk.telemetry.RegressionPrediction]s
             batch_size: Batch size for processing Dataset inputs (only used when input is a Dataset)
+            prompt: Optional prompt for instruction-tuned embedding models
             background: Whether to run the operation in the background and return a job handle
 
         Returns:
@@ -554,6 +553,14 @@ class RegressionModel:
                 r2: 0.8500,
                 anomaly_score: 0.3500 ± 0.0500,
             })
+
+            >>> # Using with an instruction-tuned embedding model
+            >>> model.evaluate(dataset,prompt="Represent this review for rating prediction:")
+            RegressionMetrics({
+            mae: 0.2000,
+            rmse: 0.3000,
+            r2: 0.9000,
+            anomaly_score: 0.3000 ± 0.0400})
         """
         if isinstance(data, Datasource):
             return self._evaluate_datasource(
@@ -572,6 +579,7 @@ class RegressionModel:
                 record_predictions=record_predictions,
                 tags=tags,
                 batch_size=batch_size,
+                prompt=prompt,
             )
         else:
             raise ValueError(f"Invalid data type: {type(data)}")
@@ -640,8 +648,9 @@ class RegressionModel:
             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)
             ],
         )

orca_sdk/regression_model_test.py

@@ -207,6 +207,21 @@ def test_predict_constraint_violation(scored_memoryset: ScoredMemoryset):
         model.predict("test")
 
 
+def test_predict_with_prompt(regression_model: RegressionModel):
+    """Test that prompt parameter is properly passed through to predictions"""
+    # Test with an instruction-supporting embedding model if available
+    prediction_with_prompt = regression_model.predict(
+        "This product is amazing!", prompt="Represent this text for rating prediction:"
+    )
+    prediction_without_prompt = regression_model.predict("This product is amazing!")
+
+    # Both should work and return valid predictions
+    assert prediction_with_prompt.score is not None
+    assert prediction_without_prompt.score is not None
+    assert 0 <= prediction_with_prompt.confidence <= 1
+    assert 0 <= prediction_without_prompt.confidence <= 1
+
+
 def test_record_prediction_feedback(regression_model: RegressionModel):
     predictions = regression_model.predict(["This is excellent!", "This is terrible!"])
     expected_scores = [0.9, 0.1]