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

orca_sdk/conftest.py

@@ -24,15 +24,6 @@ os.environ["ORCA_API_URL"] = os.environ.get("ORCA_API_URL", "http://localhost:15
 os.environ["ORCA_SAVE_TELEMETRY_SYNCHRONOUSLY"] = "true"
 
 
-def skip_in_prod(reason: str):
-    """Custom decorator to skip tests when running against production API"""
-    PROD_API_URLs = ["https://api.orcadb.ai", "https://api.staging.orcadb.ai"]
-    return pytest.mark.skipif(
-        os.environ["ORCA_API_URL"] in PROD_API_URLs,
-        reason=reason,
-    )
-
-
 def skip_in_ci(reason: str):
     """Custom decorator to skip tests when running in CI"""
     return pytest.mark.skipif(
@@ -201,6 +192,11 @@ SAMPLE_DATA = [
 ]
 
 
+@pytest.fixture(scope="session")
+def data() -> list[dict]:
+    return SAMPLE_DATA
+
+
 @pytest.fixture(scope="session")
 def hf_dataset(label_names: list[str]) -> Dataset:
     return Dataset.from_list(
@@ -232,6 +228,11 @@ EVAL_DATASET = [
 ]
 
 
+@pytest.fixture(scope="session")
+def eval_data() -> list[dict]:
+    return EVAL_DATASET
+
+
 @pytest.fixture(scope="session")
 def eval_datasource() -> Datasource:
     eval_datasource = Datasource.from_list("eval_datasource", EVAL_DATASET)
@@ -288,6 +289,7 @@ def writable_memoryset(datasource: Datasource, api_key: str) -> Generator[Labele
         datasource=datasource,
         embedding_model=PretrainedEmbeddingModel.GTE_BASE,
         source_id_column="source_id",
+        partition_id_column="partition_id",
         max_seq_length_override=32,
         if_exists="open",
     )
@@ -297,13 +299,7 @@ def writable_memoryset(datasource: Datasource, api_key: str) -> Generator[Labele
         # Restore the memoryset to a clean state for the next test.
         with OrcaClient(api_key=api_key).use():
             if LabeledMemoryset.exists("test_writable_memoryset"):
-                memoryset.refresh()
-
-                memory_ids = [memoryset[i].memory_id for i in range(len(memoryset))]
-
-                if memory_ids:
-                    memoryset.delete(memory_ids)
-                memoryset.refresh()
+                memoryset.truncate()
                 assert len(memoryset) == 0
                 memoryset.insert(SAMPLE_DATA)
         # If the test dropped the memoryset, do nothing — it will be recreated on the next use.
@@ -380,3 +376,88 @@ def partitioned_regression_model(readonly_partitioned_scored_memoryset: ScoredMe
         description="test_partitioned_regression_description",
     )
     return model
+
+
+@pytest.fixture(scope="function")
+def fully_partitioned_classification_resources() -> (
+    Generator[tuple[Datasource, LabeledMemoryset, ClassificationModel], None, None]
+):
+    data = [
+        {"value": "i love soup", "label": 0, "partition_id": "p1"},
+        {"value": "cats are cute", "label": 1, "partition_id": "p1"},
+        {"value": "soup is good", "label": 0, "partition_id": "p1"},
+        {"value": "i love cats", "label": 1, "partition_id": "p2"},
+        {"value": "everyone loves cats", "label": 1, "partition_id": "p2"},
+        {"value": "soup is good", "label": 0, "partition_id": "p1"},
+        {"value": "cats are amazing animals", "label": 1, "partition_id": "p2"},
+        {"value": "tomato soup is delicious", "label": 0, "partition_id": "p1"},
+        {"value": "cats love to play", "label": 1, "partition_id": "p2"},
+        {"value": "i enjoy eating soup", "label": 0, "partition_id": "p1"},
+        {"value": "my cat is fluffy", "label": 1, "partition_id": "p2"},
+        {"value": "chicken soup is tasty", "label": 0, "partition_id": "p1"},
+        {"value": "cats are playful pets", "label": 1, "partition_id": "p2"},
+        {"value": "soup warms the soul", "label": 0, "partition_id": "p1"},
+        {"value": "cats have soft fur", "label": 1, "partition_id": "p2"},
+        {"value": "vegetable soup is healthy", "label": 0, "partition_id": "p1"},
+    ]
+
+    datasource = None
+    memoryset = None
+    classification_model = None
+    try:
+        datasource = Datasource.from_list("fully_partitioned_classification_datasource", data)
+        memoryset = LabeledMemoryset.create(
+            "fully_partitioned_classification_memoryset",
+            datasource=datasource,
+            label_names=["soup", "cats"],
+            partition_id_column="partition_id",
+        )
+        classification_model = ClassificationModel.create("fully_partitioned_classification_model", memoryset=memoryset)
+        yield (datasource, memoryset, classification_model)
+    finally:
+        # Clean up in reverse order of creation
+        ClassificationModel.drop("fully_partitioned_classification_model", if_not_exists="ignore")
+        LabeledMemoryset.drop("fully_partitioned_classification_memoryset", if_not_exists="ignore")
+        Datasource.drop("fully_partitioned_classification_datasource", if_not_exists="ignore")
+
+
+@pytest.fixture(scope="function")
+def fully_partitioned_regression_resources() -> (
+    Generator[tuple[Datasource, ScoredMemoryset, RegressionModel], None, None]
+):
+    data = [
+        {"value": "i love soup", "score": 0.1, "partition_id": "p1"},
+        {"value": "cats are cute", "score": 0.9, "partition_id": "p1"},
+        {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
+        {"value": "i love cats", "score": 0.9, "partition_id": "p2"},
+        {"value": "everyone loves cats", "score": 0.9, "partition_id": "p2"},
+        {"value": "soup is good", "score": 0.1, "partition_id": "p1"},
+        {"value": "cats are amazing animals", "score": 0.9, "partition_id": "p2"},
+        {"value": "tomato soup is delicious", "score": 0.1, "partition_id": "p1"},
+        {"value": "cats love to play", "score": 0.9, "partition_id": "p2"},
+        {"value": "i enjoy eating soup", "score": 0.1, "partition_id": "p1"},
+        {"value": "my cat is fluffy", "score": 0.9, "partition_id": "p2"},
+        {"value": "chicken soup is tasty", "score": 0.1, "partition_id": "p1"},
+        {"value": "cats are playful pets", "score": 0.9, "partition_id": "p2"},
+        {"value": "soup warms the soul", "score": 0.1, "partition_id": "p1"},
+        {"value": "cats have soft fur", "score": 0.9, "partition_id": "p2"},
+        {"value": "vegetable soup is healthy", "score": 0.1, "partition_id": "p1"},
+    ]
+
+    datasource = None
+    memoryset = None
+    regression_model = None
+    try:
+        datasource = Datasource.from_list("fully_partitioned_regression_datasource", data)
+        memoryset = ScoredMemoryset.create(
+            "fully_partitioned_regression_memoryset",
+            datasource=datasource,
+            partition_id_column="partition_id",
+        )
+        regression_model = RegressionModel.create("fully_partitioned_regression_model", memoryset=memoryset)
+        yield (datasource, memoryset, regression_model)
+    finally:
+        # Clean up in reverse order of creation
+        RegressionModel.drop("fully_partitioned_regression_model", if_not_exists="ignore")
+        ScoredMemoryset.drop("fully_partitioned_regression_memoryset", if_not_exists="ignore")
+        Datasource.drop("fully_partitioned_regression_datasource", if_not_exists="ignore")

orca_sdk/credentials_test.py

@@ -75,7 +75,7 @@ def test_create_api_key_already_exists():
         OrcaCredentials.create_api_key("orca_sdk_test")
 
 
-def test_set_api_key(api_key):
+def test_use_client(api_key):
     client = OrcaClient(api_key=str(uuid4()))
     with client.use():
         assert not OrcaCredentials.is_authenticated()
@@ -91,17 +91,14 @@ def test_set_base_url(api_key):
     assert client.base_url == "http://localhost:1583"
 
 
-# deprecated methods:
-
-
-def test_deprecated_set_api_key(api_key):
+def test_set_api_key(api_key):
     with OrcaClient(api_key=str(uuid4())).use():
         assert not OrcaCredentials.is_authenticated()
         OrcaCredentials.set_api_key(api_key)
         assert OrcaCredentials.is_authenticated()
 
 
-def test_deprecated_set_invalid_api_key(api_key):
+def test_set_invalid_api_key(api_key):
     with OrcaClient(api_key=api_key).use():
         assert OrcaCredentials.is_authenticated()
         with pytest.raises(ValueError, match="Invalid API key"):
@@ -109,13 +106,13 @@ def test_deprecated_set_invalid_api_key(api_key):
         assert not OrcaCredentials.is_authenticated()
 
 
-def test_deprecated_set_api_url(api_key):
+def test_set_api_url(api_key):
     with OrcaClient(api_key=api_key).use():
         OrcaCredentials.set_api_url("http://api.orcadb.ai")
         assert str(OrcaClient._resolve_client().base_url) == "http://api.orcadb.ai"
 
 
-def test_deprecated_set_invalid_api_url(api_key):
+def test_set_invalid_api_url(api_key):
     with OrcaClient(api_key=api_key).use():
         with pytest.raises(ValueError, match="No API found at http://localhost:1582"):
             OrcaCredentials.set_api_url("http://localhost:1582")

orca_sdk/datasource.py

@@ -1,28 +1,30 @@
 from __future__ import annotations
 
-import logging
 import tempfile
 import zipfile
 from datetime import datetime
 from io import BytesIO
 from os import PathLike
 from pathlib import Path
-from typing import Any, Literal, Union, cast
+from typing import TYPE_CHECKING, Any, Literal, Union, cast
 
-import pandas as pd
-import pyarrow as pa
-from datasets import Dataset, DatasetDict
 from httpx._types import FileTypes  # type: ignore
-from pyarrow import parquet
-from torch.utils.data import DataLoader as TorchDataLoader
-from torch.utils.data import Dataset as TorchDataset
 from tqdm.auto import tqdm
 
-from ._utils.common import CreateMode, DropMode
-from ._utils.data_parsing import hf_dataset_from_torch
+from ._utils.common import CreateMode, DropMode, logger
+from ._utils.torch_parsing import list_from_torch
 from ._utils.tqdm_file_reader import TqdmFileReader
 from .client import DatasourceMetadata, OrcaClient
 
+if TYPE_CHECKING:
+    # These are peer dependencies that are used for types only
+    from datasets import Dataset as HFDataset  # type: ignore
+    from datasets import DatasetDict as HFDatasetDict  # type: ignore
+    from pandas import DataFrame as PandasDataFrame  # type: ignore
+    from pyarrow import Table as PyArrowTable  # type: ignore
+    from torch.utils.data import DataLoader as TorchDataLoader  # type: ignore
+    from torch.utils.data import Dataset as TorchDataset  # type: ignore
+
 
 def _upload_files_to_datasource(
     name: str,
@@ -144,7 +146,7 @@ class Datasource:
 
     @classmethod
     def from_hf_dataset(
-        cls, name: str, dataset: Dataset, if_exists: CreateMode = "error", description: str | None = None
+        cls, name: str, dataset: HFDataset, if_exists: CreateMode = "error", description: str | None = None
     ) -> Datasource:
         """
         Create a new datasource from a Hugging Face Dataset
@@ -181,7 +183,7 @@ class Datasource:
     def from_hf_dataset_dict(
         cls,
         name: str,
-        dataset_dict: DatasetDict,
+        dataset_dict: HFDatasetDict,
         if_exists: CreateMode = "error",
         description: dict[str, str | None] | str | None = None,
     ) -> dict[str, Datasource]:
@@ -237,8 +239,8 @@ class Datasource:
         Raises:
             ValueError: If the datasource already exists and if_exists is `"error"`
         """
-        hf_dataset = hf_dataset_from_torch(torch_data, column_names=column_names)
-        return cls.from_hf_dataset(name, hf_dataset, if_exists=if_exists, description=description)
+        data_list = list_from_torch(torch_data, column_names=column_names)
+        return cls.from_list(name, data_list, if_exists=if_exists, description=description)
 
     @classmethod
     def from_list(
@@ -312,7 +314,7 @@ class Datasource:
 
     @classmethod
     def from_pandas(
-        cls, name: str, dataframe: pd.DataFrame, if_exists: CreateMode = "error", description: str | None = None
+        cls, name: str, dataframe: PandasDataFrame, if_exists: CreateMode = "error", description: str | None = None
     ) -> Datasource:
         """
         Create a new datasource from a pandas DataFrame
@@ -324,18 +326,28 @@ class Datasource:
                 `"error"`. Other option is `"open"` to open the existing datasource.
             description: Optional description for the datasource
 
+        Notes:
+            Data type precision may be lost during upload unless the [`datasets`][datasets] library is installed.
+
         Returns:
             A handle to the new datasource in the OrcaCloud
 
         Raises:
             ValueError: If the datasource already exists and if_exists is `"error"`
+            ImportError: If the upload dependency group is not installed
         """
-        dataset = Dataset.from_pandas(dataframe)
-        return cls.from_hf_dataset(name, dataset, if_exists=if_exists, description=description)
+        try:
+            from datasets import Dataset  # type: ignore
+
+            return cls.from_hf_dataset(
+                name, Dataset.from_pandas(dataframe), if_exists=if_exists, description=description
+            )
+        except ImportError:
+            return cls.from_dict(name, dataframe.to_dict(orient="list"), if_exists=if_exists, description=description)
 
     @classmethod
     def from_arrow(
-        cls, name: str, pyarrow_table: pa.Table, if_exists: CreateMode = "error", description: str | None = None
+        cls, name: str, pyarrow_table: PyArrowTable, if_exists: CreateMode = "error", description: str | None = None
     ) -> Datasource:
         """
         Create a new datasource from a pyarrow Table
@@ -358,6 +370,9 @@ class Datasource:
         if existing is not None:
             return existing
 
+        # peer dependency that is guaranteed to exist if the user provided a pyarrow table
+        from pyarrow import parquet  # type: ignore
+
         # Write to bytes buffer
         buffer = BytesIO()
         parquet.write_table(pyarrow_table, buffer)
@@ -399,6 +414,7 @@ class Datasource:
 
         Raises:
             ValueError: If the datasource already exists and if_exists is `"error"`
+            ImportError: If the path is a directory and [`datasets`][datasets] is not installed
         """
         # Check if datasource already exists and handle accordingly
         existing = _handle_existing_datasource(name, if_exists)
@@ -409,6 +425,13 @@ class Datasource:
 
         # For dataset directories, use the upload endpoint with multiple files
         if file_path.is_dir():
+            try:
+                from datasets import Dataset  # type: ignore
+            except ImportError as e:
+                raise ImportError(
+                    "The path is a directory, we only support uploading directories that contain saved HuggingFace datasets but datasets is not installed."
+                ) from e
+
             return cls.from_hf_dataset(
                 name, Dataset.load_from_disk(file_path), if_exists=if_exists, description=description
             )
@@ -479,7 +502,7 @@ class Datasource:
         try:
             client = OrcaClient._resolve_client()
             client.DELETE("/datasource/{name_or_id}", params={"name_or_id": name_or_id})
-            logging.info(f"Deleted datasource {name_or_id}")
+            logger.info(f"Deleted datasource {name_or_id}")
         except LookupError:
             if if_not_exists == "error":
                 raise
@@ -561,9 +584,9 @@ class Datasource:
             with zipfile.ZipFile(output_path, "r") as zip_ref:
                 zip_ref.extractall(extract_dir)
             output_path.unlink()  # Remove the zip file after extraction
-            logging.info(f"Downloaded {extract_dir}")
+            logger.info(f"Downloaded {extract_dir}")
         else:
-            logging.info(f"Downloaded {output_path}")
+            logger.info(f"Downloaded {output_path}")
 
     def to_list(self) -> list[dict]:
         """

orca_sdk/datasource_test.py

@@ -5,8 +5,6 @@ from typing import cast
 from uuid import uuid4
 
 import numpy as np
-import pandas as pd
-import pyarrow as pa
 import pytest
 from datasets import Dataset
 
@@ -137,6 +135,8 @@ def test_from_dict():
 
 
 def test_from_pandas():
+    pd = pytest.importorskip("pandas")
+
     # Test creating datasource from pandas DataFrame
     df = pd.DataFrame(
         {
@@ -152,6 +152,8 @@ def test_from_pandas():
 
 
 def test_from_arrow():
+    pa = pytest.importorskip("pyarrow")
+
     # Test creating datasource from pyarrow Table
     table = pa.table(
         {
@@ -205,6 +207,8 @@ def test_from_dict_already_exists():
 
 
 def test_from_pandas_already_exists():
+    pd = pytest.importorskip("pandas")
+
     # Test the if_exists parameter with from_pandas
     df = pd.DataFrame({"column1": [1], "column2": ["a"]})
     name = f"test_pandas_exists_{uuid4()}"
@@ -224,6 +228,8 @@ def test_from_pandas_already_exists():
 
 
 def test_from_arrow_already_exists():
+    pa = pytest.importorskip("pyarrow")
+
     # Test the if_exists parameter with from_arrow
     table = pa.table({"column1": [1], "column2": ["a"]})
     name = f"test_arrow_exists_{uuid4()}"

orca_sdk/embedding_model.py

@@ -4,8 +4,7 @@ from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import TYPE_CHECKING, Literal, Sequence, cast, get_args, overload
 
-from ._shared.metrics import ClassificationMetrics, RegressionMetrics
-from ._utils.common import UNSET, CreateMode, DropMode
+from ._utils.common import CreateMode, DropMode
 from .client import (
     EmbeddingEvaluationRequest,
     EmbeddingFinetuningMethod,
@@ -20,7 +19,9 @@ from .datasource import Datasource
 from .job import Job, Status
 
 if TYPE_CHECKING:
+    from .classification_model import ClassificationMetrics
     from .memoryset import LabeledMemoryset, ScoredMemoryset
+    from .regression_model import RegressionMetrics
 
 
 class EmbeddingModelBase(ABC):
@@ -230,6 +231,9 @@ class EmbeddingModelBase(ABC):
             raise ValueError("Invalid embedding model")
 
         def get_result(job_id: str) -> ClassificationMetrics | RegressionMetrics:
+            from .classification_model import ClassificationMetrics
+            from .regression_model import RegressionMetrics
+
             client = OrcaClient._resolve_client()
             if isinstance(self, PretrainedEmbeddingModel):
                 res = client.GET(
@@ -244,34 +248,7 @@ class EmbeddingModelBase(ABC):
             else:
                 raise ValueError("Invalid embedding model")
             assert res is not None
-            return (
-                RegressionMetrics(
-                    coverage=res.get("coverage"),
-                    mse=res.get("mse"),
-                    rmse=res.get("rmse"),
-                    mae=res.get("mae"),
-                    r2=res.get("r2"),
-                    explained_variance=res.get("explained_variance"),
-                    loss=res.get("loss"),
-                    anomaly_score_mean=res.get("anomaly_score_mean"),
-                    anomaly_score_median=res.get("anomaly_score_median"),
-                    anomaly_score_variance=res.get("anomaly_score_variance"),
-                )
-                if "mse" in res
-                else ClassificationMetrics(
-                    coverage=res.get("coverage"),
-                    f1_score=res.get("f1_score"),
-                    accuracy=res.get("accuracy"),
-                    loss=res.get("loss"),
-                    anomaly_score_mean=res.get("anomaly_score_mean"),
-                    anomaly_score_median=res.get("anomaly_score_median"),
-                    anomaly_score_variance=res.get("anomaly_score_variance"),
-                    roc_auc=res.get("roc_auc"),
-                    pr_auc=res.get("pr_auc"),
-                    pr_curve=res.get("pr_curve"),
-                    roc_curve=res.get("roc_curve"),
-                )
-            )
+            return RegressionMetrics(res) if "mse" in res else ClassificationMetrics(res)
 
         job = Job(response["job_id"], lambda: get_result(response["job_id"]))
         return job if background else job.result()
@@ -404,7 +381,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         return isinstance(other, PretrainedEmbeddingModel) and self.name == other.name
 
     def __repr__(self) -> str:
-        return f"PretrainedEmbeddingModel({{name: {self.name}, embedding_dim: {self.embedding_dim}, max_seq_length: {self.max_seq_length}, num_params: {self.num_params/1000000:.0f}M}})"
+        return f"PretrainedEmbeddingModel({{name: {self.name}, embedding_dim: {self.embedding_dim}, max_seq_length: {self.max_seq_length}, num_params: {self.num_params / 1000000:.0f}M}})"
 
     @classmethod
     def all(cls) -> list[PretrainedEmbeddingModel]:
@@ -691,21 +668,26 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
             return False
 
     @classmethod
-    def drop(cls, name_or_id: str, *, if_not_exists: DropMode = "error"):
+    def drop(cls, name_or_id: str, *, if_not_exists: DropMode = "error", cascade: bool = False):
         """
         Delete the finetuned embedding model from the OrcaCloud
 
         Params:
             name_or_id: The name or id of the finetuned embedding model
+            if_not_exists: What to do if the finetuned embedding model does not exist, defaults to `"error"`.
+                Other option is `"ignore"` to do nothing if the model does not exist.
+            cascade: If True, also delete all associated memorysets and their predictive models.
+                Defaults to False.
 
         Raises:
             LookupError: If the finetuned embedding model does not exist and `if_not_exists` is `"error"`
+            RuntimeError: If the model has associated memorysets and cascade is False
         """
         try:
             client = OrcaClient._resolve_client()
             client.DELETE(
                 "/finetuned_embedding_model/{name_or_id}",
-                params={"name_or_id": name_or_id},
+                params={"name_or_id": name_or_id, "cascade": cascade},
             )
         except LookupError:
             if if_not_exists == "error":

orca_sdk/embedding_model_test.py

@@ -4,9 +4,9 @@ from uuid import uuid4
 
 import pytest
 
+from .classification_model import ClassificationMetrics
 from .datasource import Datasource
 from .embedding_model import (
-    ClassificationMetrics,
     FinetunedEmbeddingModel,
     PretrainedEmbeddingModel,
     PretrainedEmbeddingModelName,
@@ -172,6 +172,35 @@ def test_drop_finetuned_model(datasource: Datasource):
         FinetunedEmbeddingModel.open("finetuned_model_to_delete")
 
 
+def test_drop_finetuned_model_with_memoryset_cascade(datasource: Datasource):
+    """Test that cascade=False prevents deletion and cascade=True allows it."""
+    finetuned_model = PretrainedEmbeddingModel.DISTILBERT.finetune("finetuned_model_cascade_delete", datasource)
+    memoryset = LabeledMemoryset.create(
+        "test_memoryset_for_finetuned_model_cascade",
+        datasource=datasource,
+        embedding_model=finetuned_model,
+    )
+
+    # Verify memoryset exists and uses the finetuned model
+    assert LabeledMemoryset.open(memoryset.name) is not None
+    assert memoryset.embedding_model == finetuned_model
+
+    # Without cascade, deletion should fail
+    with pytest.raises(RuntimeError):
+        FinetunedEmbeddingModel.drop(finetuned_model.id, cascade=False)
+
+    # Model and memoryset should still exist
+    assert FinetunedEmbeddingModel.exists(finetuned_model.name)
+    assert LabeledMemoryset.exists(memoryset.name)
+
+    # With cascade, deletion should succeed
+    FinetunedEmbeddingModel.drop(finetuned_model.id, cascade=True)
+
+    # Both model and memoryset should be deleted
+    assert not FinetunedEmbeddingModel.exists(finetuned_model.name)
+    assert not LabeledMemoryset.exists(memoryset.name)
+
+
 def test_drop_finetuned_model_unauthenticated(unauthenticated_client, datasource: Datasource):
     with unauthenticated_client.use():
         with pytest.raises(ValueError, match="Invalid API key"):