orca-sdk 0.1.9__py3-none-any.whl

orca_sdk/conftest.py

@@ -0,0 +1,382 @@
+import logging
+import os
+from typing import Generator
+from uuid import uuid4
+
+import pytest
+import pytest_asyncio
+from datasets import ClassLabel, Dataset, Features, Value
+
+from ._utils.auth import _create_api_key, _delete_org
+from .async_client import OrcaAsyncClient
+from .classification_model import ClassificationModel
+from .client import OrcaClient
+from .credentials import OrcaCredentials
+from .datasource import Datasource
+from .embedding_model import PretrainedEmbeddingModel
+from .memoryset import LabeledMemoryset, ScoredMemoryset
+from .regression_model import RegressionModel
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+
+os.environ["ORCA_API_URL"] = os.environ.get("ORCA_API_URL", "http://localhost:1584/")
+
+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(
+        os.environ.get("GITHUB_ACTIONS", "false") == "true",
+        reason=reason,
+    )
+
+
+def _create_org_id():
+    # UUID start to identify test data (0xtest...)
+    return "10e50000-0000-4000-a000-" + str(uuid4())[24:]
+
+
+@pytest.fixture(scope="session")
+def org_id():
+    return _create_org_id()
+
+
+@pytest.fixture(scope="session")
+def other_org_id():
+    return _create_org_id()
+
+
+@pytest.fixture(autouse=True, scope="session")
+def api_key(org_id) -> Generator[str, None, None]:
+    api_key = _create_api_key(org_id=org_id, name="orca_sdk_test")
+    with OrcaClient(api_key=api_key).use():
+        yield api_key
+    _delete_org(org_id)
+
+
+# We cannot use a session scoped fixture because async pytest tears down the client after each test
+@pytest.fixture(autouse=True)
+def authenticate_async_client(api_key) -> Generator[None, None, None]:
+    with OrcaAsyncClient(api_key=api_key).use():
+        yield
+
+
+@pytest.fixture(scope="session")
+def unauthenticated_client() -> OrcaClient:
+    return OrcaClient(api_key=str(uuid4()))
+
+
+@pytest_asyncio.fixture()
+def unauthenticated_async_client() -> OrcaAsyncClient:
+    return OrcaAsyncClient(api_key=str(uuid4()))
+
+
+@pytest.fixture(scope="session")
+def unauthorized_client(other_org_id):
+    different_api_key = _create_api_key(org_id=other_org_id, name="orca_sdk_test_other_org")
+    return OrcaClient(api_key=different_api_key)
+
+
+@pytest.fixture(scope="session")
+def predict_only_client() -> OrcaClient:
+    predict_api_key = OrcaCredentials.create_api_key("orca_sdk_test_predict", scopes={"PREDICT"})
+    return OrcaClient(api_key=predict_api_key)
+
+
+@pytest.fixture(scope="session")
+def label_names():
+    return ["soup", "cats"]
+
+
+SAMPLE_DATA = [
+    {"value": "i love soup", "label": 0, "key": "g1", "score": 0.1, "source_id": "s1", "partition_id": "p1"},
+    {"value": "cats are cute", "label": 1, "key": "g1", "score": 0.9, "source_id": "s2", "partition_id": "p1"},
+    {"value": "soup is good", "label": 0, "key": "g1", "score": 0.1, "source_id": "s3", "partition_id": "p1"},
+    {"value": "i love cats", "label": 1, "key": "g1", "score": 0.9, "source_id": "s4", "partition_id": "p1"},
+    {"value": "everyone loves cats", "label": 1, "key": "g1", "score": 0.9, "source_id": "s5", "partition_id": "p1"},
+    {
+        "value": "soup is great for the winter",
+        "label": 0,
+        "key": "g1",
+        "score": 0.1,
+        "source_id": "s6",
+        "partition_id": "p1",
+    },
+    {
+        "value": "hot soup on a rainy day!",
+        "label": 0,
+        "key": "g1",
+        "score": 0.1,
+        "source_id": "s7",
+        "partition_id": "p1",
+    },
+    {"value": "cats sleep all day", "label": 1, "key": "g1", "score": 0.9, "source_id": "s8", "partition_id": "p1"},
+    {"value": "homemade soup recipes", "label": 0, "key": "g1", "score": 0.1, "source_id": "s9", "partition_id": "p2"},
+    {"value": "cats purr when happy", "label": 1, "key": "g2", "score": 0.9, "source_id": "s10", "partition_id": "p2"},
+    {
+        "value": "chicken noodle soup is classic",
+        "label": 0,
+        "key": "g1",
+        "score": 0.1,
+        "source_id": "s11",
+        "partition_id": "p2",
+    },
+    {"value": "kittens are baby cats", "label": 1, "key": "g2", "score": 0.9, "source_id": "s12", "partition_id": "p2"},
+    {
+        "value": "soup can be served cold too",
+        "label": 0,
+        "key": "g1",
+        "score": 0.1,
+        "source_id": "s13",
+        "partition_id": "p2",
+    },
+    {"value": "cats have nine lives", "label": 1, "key": "g2", "score": 0.9, "source_id": "s14", "partition_id": "p2"},
+    {
+        "value": "tomato soup with grilled cheese",
+        "label": 0,
+        "key": "g1",
+        "score": 0.1,
+        "source_id": "s15",
+        "partition_id": "p2",
+    },
+    {
+        "value": "cats are independent animals",
+        "label": 1,
+        "key": "g2",
+        "score": 0.9,
+        "source_id": "s16",
+        "partition_id": None,
+    },
+    {
+        "value": "the beach is always fun",
+        "label": None,
+        "key": "g3",
+        "score": None,
+        "source_id": "s17",
+        "partition_id": None,
+    },
+    {"value": "i love the beach", "label": None, "key": "g3", "score": None, "source_id": "s18", "partition_id": None},
+    {
+        "value": "the ocean is healing",
+        "label": None,
+        "key": "g3",
+        "score": None,
+        "source_id": "s19",
+        "partition_id": None,
+    },
+    {
+        "value": "sandy feet, sand between my toes at the beach",
+        "label": None,
+        "key": "g3",
+        "score": None,
+        "source_id": "s20",
+        "partition_id": None,
+    },
+    {
+        "value": "i am such a beach bum",
+        "label": None,
+        "key": "g3",
+        "score": None,
+        "source_id": "s21",
+        "partition_id": None,
+    },
+    {
+        "value": "i will always want to be at the beach",
+        "label": None,
+        "key": "g3",
+        "score": None,
+        "source_id": "s22",
+        "partition_id": None,
+    },
+]
+
+
+@pytest.fixture(scope="session")
+def hf_dataset(label_names: list[str]) -> Dataset:
+    return Dataset.from_list(
+        SAMPLE_DATA,
+        features=Features(
+            {
+                "value": Value("string"),
+                "label": ClassLabel(names=label_names),
+                "key": Value("string"),
+                "score": Value("float"),
+                "source_id": Value("string"),
+                "partition_id": Value("string"),
+            }
+        ),
+    )
+
+
+@pytest.fixture(scope="session")
+def datasource(hf_dataset: Dataset) -> Datasource:
+    datasource = Datasource.from_hf_dataset("test_datasource", hf_dataset)
+    return datasource
+
+
+EVAL_DATASET = [
+    {"value": "chicken noodle soup is the best", "label": 1, "score": 0.9},  # mislabeled
+    {"value": "cats are cute", "label": 0, "score": 0.1},  # mislabeled
+    {"value": "soup is great for the winter", "label": 0, "score": 0.1},
+    {"value": "i love cats", "label": 1, "score": 0.9},
+]
+
+
+@pytest.fixture(scope="session")
+def eval_datasource() -> Datasource:
+    eval_datasource = Datasource.from_list("eval_datasource", EVAL_DATASET)
+    return eval_datasource
+
+
+@pytest.fixture(scope="session")
+def eval_dataset() -> Dataset:
+    eval_dataset = Dataset.from_list(EVAL_DATASET)
+    return eval_dataset
+
+
+@pytest.fixture(scope="session")
+def readonly_memoryset(datasource: Datasource) -> LabeledMemoryset:
+    memoryset = LabeledMemoryset.create(
+        "test_readonly_memoryset",
+        datasource=datasource,
+        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
+        source_id_column="source_id",
+        max_seq_length_override=32,
+        index_type="IVF_FLAT",
+        index_params={"n_lists": 100},
+    )
+    return memoryset
+
+
+@pytest.fixture(scope="session")
+def readonly_partitioned_memoryset(datasource: Datasource) -> LabeledMemoryset:
+    memoryset = LabeledMemoryset.create(
+        "test_readonly_partitioned_memoryset",
+        datasource=datasource,
+        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
+        source_id_column="source_id",
+        partition_id_column="partition_id",
+    )
+    return memoryset
+
+
+@pytest.fixture(scope="function")
+def writable_memoryset(datasource: Datasource, api_key: str) -> Generator[LabeledMemoryset, None, None]:
+    """
+    Function-scoped fixture that provides a writable memoryset for tests that mutate state.
+
+    This fixture creates a fresh `LabeledMemoryset` named 'test_writable_memoryset' before each test.
+    After the test, it attempts to restore the memoryset to its initial state by deleting any added entries
+    and reinserting sample data — unless the memoryset has been dropped by the test itself, in which case
+    it will be recreated on the next invocation.
+
+    Note: Re-creating the memoryset from scratch is surprisingly more expensive than cleaning it up.
+    """
+    # It shouldn't be possible for this memoryset to already exist
+    memoryset = LabeledMemoryset.create(
+        "test_writable_memoryset",
+        datasource=datasource,
+        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
+        source_id_column="source_id",
+        max_seq_length_override=32,
+        if_exists="open",
+    )
+    try:
+        yield memoryset
+    finally:
+        # 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()
+                assert len(memoryset) == 0
+                memoryset.insert(SAMPLE_DATA)
+        # If the test dropped the memoryset, do nothing — it will be recreated on the next use.
+
+
+@pytest.fixture(scope="session")
+def classification_model(readonly_memoryset: LabeledMemoryset) -> ClassificationModel:
+    model = ClassificationModel.create(
+        "test_classification_model",
+        readonly_memoryset,
+        num_classes=2,
+        memory_lookup_count=3,
+        description="test_description",
+    )
+    return model
+
+
+@pytest.fixture(scope="session")
+def partitioned_classification_model(readonly_partitioned_memoryset: LabeledMemoryset) -> ClassificationModel:
+    model = ClassificationModel.create(
+        "test_partitioned_classification_model",
+        readonly_partitioned_memoryset,
+        num_classes=2,
+        memory_lookup_count=3,
+        description="test_partitioned_description",
+    )
+    return model
+
+
+# Add scored memoryset and regression model fixtures
+@pytest.fixture(scope="session")
+def scored_memoryset(datasource: Datasource) -> ScoredMemoryset:
+    memoryset = ScoredMemoryset.create(
+        "test_scored_memoryset",
+        datasource=datasource,
+        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
+        source_id_column="source_id",
+        max_seq_length_override=32,
+        index_type="IVF_FLAT",
+        index_params={"n_lists": 100},
+    )
+    return memoryset
+
+
+@pytest.fixture(scope="session")
+def regression_model(scored_memoryset: ScoredMemoryset) -> RegressionModel:
+    model = RegressionModel.create(
+        "test_regression_model",
+        scored_memoryset,
+        memory_lookup_count=3,
+        description="test_regression_description",
+    )
+    return model
+
+
+@pytest.fixture(scope="session")
+def readonly_partitioned_scored_memoryset(datasource: Datasource) -> ScoredMemoryset:
+    memoryset = ScoredMemoryset.create(
+        "test_readonly_partitioned_scored_memoryset",
+        datasource=datasource,
+        embedding_model=PretrainedEmbeddingModel.GTE_BASE,
+        source_id_column="source_id",
+        partition_id_column="partition_id",
+    )
+    return memoryset
+
+
+@pytest.fixture(scope="session")
+def partitioned_regression_model(readonly_partitioned_scored_memoryset: ScoredMemoryset) -> RegressionModel:
+    model = RegressionModel.create(
+        "test_partitioned_regression_model",
+        readonly_partitioned_scored_memoryset,
+        memory_lookup_count=3,
+        description="test_partitioned_regression_description",
+    )
+    return model

orca_sdk/credentials.py

@@ -0,0 +1,217 @@
+import os
+from datetime import datetime
+from typing import Literal, NamedTuple
+
+import httpx
+from httpx import ConnectError, Headers, HTTPTransport
+from typing_extensions import deprecated
+
+from .async_client import OrcaAsyncClient
+from .client import OrcaClient
+
+Scope = Literal["ADMINISTER", "PREDICT"]
+"""
+The scopes of an API key.
+
+- `ADMINISTER`: Can do anything, including creating and deleting organizations, models, and API keys.
+- `PREDICT`: Can only call model.predict and perform CRUD operations on predictions.
+"""
+
+
+class ApiKeyInfo:
+    """
+    Information about an API key
+
+    Note:
+        The value of the API key is only available at creation time.
+
+    Attributes:
+        name: Unique name of the API key
+        created_at: When the API key was created
+        scopes: The scopes of the API key
+    """
+
+    name: str
+    created_at: datetime
+    scopes: set[Scope]
+
+    def __init__(self, name: str, created_at: datetime, scopes: set[Scope]):
+        self.name = name
+        self.created_at = created_at
+        self.scopes = scopes
+
+    def __repr__(self) -> str:
+        return "ApiKey({ " + f"name: '{self.name}', scopes: <{'|'.join(self.scopes)}>" + "})"
+
+
+class OrcaCredentials:
+    """
+    Class for managing Orca API credentials
+    """
+
+    @staticmethod
+    def is_authenticated() -> bool:
+        """
+        Check if you are authenticated to interact with the Orca API
+
+        Returns:
+            True if you are authenticated, False otherwise
+        """
+        client = OrcaClient._resolve_client()
+        try:
+            return client.GET("/auth")
+        except ValueError as e:
+            if "Invalid API key" in str(e):
+                return False
+            raise e
+
+    @staticmethod
+    def is_healthy() -> bool:
+        """
+        Check whether the API is healthy
+
+        Returns:
+            True if the API is healthy, False otherwise
+        """
+        client = OrcaClient._resolve_client()
+        try:
+            # we don't want a retry transport here, so we use httpx directly
+            httpx.get(f"{client.base_url}/check/healthy")
+        except Exception:
+            return False
+        return True
+
+    @staticmethod
+    def list_api_keys() -> list[ApiKeyInfo]:
+        """
+        List all API keys that have been created for your org
+
+        Returns:
+            A list of named tuples, with the name and creation date time of the API key
+        """
+        client = OrcaClient._resolve_client()
+        return [
+            ApiKeyInfo(
+                name=api_key["name"],
+                created_at=datetime.fromisoformat(api_key["created_at"]),
+                scopes=set(api_key["scope"]),
+            )
+            for api_key in client.GET("/auth/api_key")
+        ]
+
+    @staticmethod
+    def create_api_key(name: str, scopes: set[Scope] = {"ADMINISTER"}) -> str:
+        """
+        Create a new API key with the given name and scopes
+
+        Params:
+            name: The name of the API key
+            scopes: The scopes of the API key
+
+        Returns:
+            The secret value of the API key. Make sure to save this value as it will not be shown again.
+        """
+        client = OrcaClient._resolve_client()
+        res = client.POST(
+            "/auth/api_key",
+            json={"name": name, "scope": list(scopes)},
+        )
+        return res["api_key"]
+
+    @staticmethod
+    def revoke_api_key(name: str) -> None:
+        """
+        Delete an API key
+
+        Params:
+            name: The name of the API key to delete
+
+        Raises:
+            ValueError: if the API key is not found
+        """
+        client = OrcaClient._resolve_client()
+        client.DELETE("/auth/api_key/{name_or_id}", params={"name_or_id": name})
+
+    # TODO: remove deprecated methods after 2026-01-01
+
+    @deprecated("Use `OrcaClient.api_key` instead")
+    @staticmethod
+    def set_api_key(api_key: str, check_validity: bool = True):
+        """
+        Set the API key to use for authenticating with the Orca API
+
+        Note:
+            The API key can also be provided by setting the `ORCA_API_KEY` environment variable
+
+        Params:
+            api_key: The API key to set
+            check_validity: Whether to check if the API key is valid and raise an error otherwise
+
+        Raises:
+            ValueError: if the API key is invalid and `check_validity` is True
+        """
+        sync_client = OrcaClient._resolve_client()
+        sync_client.api_key = api_key
+        if check_validity:
+            sync_client.GET("/auth")
+
+        async_client = OrcaAsyncClient._resolve_client()
+        async_client.api_key = api_key
+
+    @deprecated("Use `OrcaClient.base_url` instead")
+    @staticmethod
+    def get_api_url() -> str:
+        """
+        Get the base URL of the Orca API that is currently being used
+        """
+        client = OrcaClient._resolve_client()
+        return str(client.base_url)
+
+    @deprecated("Use `OrcaClient.base_url` instead")
+    @staticmethod
+    def set_api_url(url: str, check_validity: bool = True):
+        """
+        Set the base URL for the Orca API
+
+        Args:
+            url: The base URL to set
+            check_validity: Whether to check if there is an API running at the given base URL
+
+        Raises:
+            ValueError: if there is no healthy API running at the given base URL and `check_validity` is True
+        """
+        # check if the base url is reachable before setting it
+        if check_validity:
+            try:
+                httpx.get(url, timeout=1)
+            except ConnectError as e:
+                raise ValueError(f"No API found at {url}") from e
+
+        sync_client = OrcaClient._resolve_client()
+        sync_client.base_url = url
+
+        async_client = OrcaAsyncClient._resolve_client()
+        async_client.base_url = url
+
+        # check if the api passes the health check
+        if check_validity:
+            OrcaCredentials.is_healthy()
+
+    @deprecated("Use `OrcaClient.headers` instead")
+    @staticmethod
+    def set_api_headers(headers: dict[str, str]):
+        """
+        Add or override default HTTP headers for all Orca API requests.
+
+        Params:
+            headers: Mapping of header names to their string values
+
+        Notes:
+            New keys are merged into the existing headers, this will overwrite headers with the
+            same name, but leave other headers untouched.
+        """
+        sync_client = OrcaClient._resolve_client()
+        sync_client.headers.update(Headers(headers))
+
+        async_client = OrcaAsyncClient._resolve_client()
+        async_client.headers.update(Headers(headers))