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

orca_sdk/conftest.py

@@ -8,6 +8,7 @@ from datasets import ClassLabel, Dataset, Features, Value
 
 from ._utils.auth import _create_api_key, _delete_org
 from .classification_model import ClassificationModel
+from .client import orca_api
 from .credentials import OrcaCredentials
 from .datasource import Datasource
 from .embedding_model import PretrainedEmbeddingModel
@@ -43,6 +44,13 @@ def _create_org_id():
     return "10e50000-0000-4000-a000-" + str(uuid4())[24:]
 
 
+@pytest.fixture()
+def base_url_reset():
+    original_base_url = orca_api.base_url
+    yield
+    orca_api.base_url = original_base_url
+
+
 @pytest.fixture(scope="session")
 def org_id():
     return _create_org_id()
@@ -192,6 +200,8 @@ def writable_memoryset(datasource: Datasource, api_key: str) -> Generator[Labele
         OrcaCredentials.set_api_key(api_key, check_validity=False)
 
         if LabeledMemoryset.exists("test_writable_memoryset"):
+            memoryset.refresh()
+
             memory_ids = [memoryset[i].memory_id for i in range(len(memoryset))]
 
             if memory_ids:

orca_sdk/credentials.py

@@ -1,17 +1,10 @@
 from datetime import datetime
 from typing import Literal, NamedTuple
 
-from ._generated_api_client.api import (
-    check_authentication,
-    create_api_key,
-    delete_api_key,
-    list_api_keys,
-)
-from ._generated_api_client.client import get_base_url, get_headers, set_headers
-from ._generated_api_client.models import (
-    CreateApiKeyRequest,
-    CreateApiKeyRequestScopeItem,
-)
+import httpx
+from httpx import ConnectError, Headers
+
+from .client import orca_api
 
 Scope = Literal["ADMINISTER", "PREDICT"]
 """
@@ -46,7 +39,7 @@ class OrcaCredentials:
         """
         Get the Orca API base URL that is currently being used
         """
-        return get_base_url()
+        return str(orca_api.base_url)
 
     @staticmethod
     def list_api_keys() -> list[ApiKeyInfo]:
@@ -57,8 +50,12 @@ class OrcaCredentials:
             A list of named tuples, with the name and creation date time of the API key
         """
         return [
-            ApiKeyInfo(name=api_key.name, created_at=api_key.created_at, scopes=set(s.value for s in api_key.scope))
-            for api_key in list_api_keys()
+            ApiKeyInfo(
+                name=api_key["name"],
+                created_at=datetime.fromisoformat(api_key["created_at"]),
+                scopes=set(api_key["scope"]),
+            )
+            for api_key in orca_api.GET("/auth/api_key")
         ]
 
     @staticmethod
@@ -70,7 +67,7 @@ class OrcaCredentials:
             True if you are authenticated, False otherwise
         """
         try:
-            return check_authentication()
+            return orca_api.GET("/auth")
         except ValueError as e:
             if "Invalid API key" in str(e):
                 return False
@@ -88,10 +85,11 @@ class OrcaCredentials:
         Returns:
             The secret value of the API key. Make sure to save this value as it will not be shown again.
         """
-        res = create_api_key(
-            body=CreateApiKeyRequest(name=name, scope=[CreateApiKeyRequestScopeItem(scope) for scope in scopes])
+        res = orca_api.POST(
+            "/auth/api_key",
+            json={"name": name, "scope": list(scopes)},
         )
-        return res.api_key
+        return res["api_key"]
 
     @staticmethod
     def revoke_api_key(name: str) -> None:
@@ -104,7 +102,7 @@ class OrcaCredentials:
         Raises:
             ValueError: if the API key is not found
         """
-        delete_api_key(name_or_id=name)
+        orca_api.DELETE("/auth/api_key/{name_or_id}", params={"name_or_id": name})
 
     @staticmethod
     def set_headers(headers: dict[str, str]):
@@ -118,7 +116,7 @@ class OrcaCredentials:
             New keys are merged into the existing headers, this will overwrite headers with the
             same name, but leave other headers untouched.
         """
-        set_headers(get_headers() | headers)
+        orca_api.headers.update(Headers(headers))
 
     @staticmethod
     def set_api_key(api_key: str, check_validity: bool = True):
@@ -137,4 +135,44 @@ class OrcaCredentials:
         """
         OrcaCredentials.set_headers({"Api-Key": api_key})
         if check_validity:
-            check_authentication()
+            orca_api.GET("/auth")
+
+    @staticmethod
+    def set_base_url(base_url: str, check_validity: bool = True):
+        """
+        Set the base URL for the Orca API
+
+        Args:
+            base_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(base_url, timeout=1)
+            except ConnectError as e:
+                raise ValueError(f"No API found at {base_url}") from e
+
+        orca_api.base_url = base_url
+
+        # check if the api passes the health check
+        if check_validity:
+            orca_api.GET("/")
+
+    @staticmethod
+    def is_healthy() -> bool:
+        """
+        Check whether the API is healthy
+
+        Returns:
+            True if the API is healthy, False otherwise
+        """
+        try:
+            orca_api.GET("/")
+            orca_api.GET("/gpu/")
+        except Exception:
+            return False
+        return True

orca_sdk/credentials_test.py

@@ -2,6 +2,7 @@ from uuid import uuid4
 
 import pytest
 
+from .client import orca_api
 from .credentials import OrcaCredentials
 
 
@@ -35,3 +36,22 @@ def test_set_invalid_api_key(api_key):
     with pytest.raises(ValueError, match="Invalid API key"):
         OrcaCredentials.set_api_key(str(uuid4()))
     assert not OrcaCredentials.is_authenticated()
+
+
+def test_set_base_url(base_url_reset):
+    OrcaCredentials.set_base_url("http://api.orcadb.ai")
+    assert str(orca_api.base_url) == "http://api.orcadb.ai"
+
+
+def test_set_invalid_base_url():
+    with pytest.raises(ValueError, match="No API found at http://localhost:1582"):
+        OrcaCredentials.set_base_url("http://localhost:1582")
+
+
+def test_is_healthy():
+    assert OrcaCredentials.is_healthy()
+
+
+def test_is_healthy_false(base_url_reset):
+    OrcaCredentials.set_base_url("http://localhost:1582", check_validity=False)
+    assert not OrcaCredentials.is_healthy()

orca_sdk/datasource.py

@@ -7,34 +7,21 @@ from datetime import datetime
 from io import BytesIO
 from os import PathLike
 from pathlib import Path
-from typing import Union
+from typing import 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 ._generated_api_client.api import (
-    create_datasource_from_content,
-    delete_datasource,
-    get_datasource,
-    list_datasources,
-)
-from ._generated_api_client.api.datasource.create_datasource_from_files_datasource_upload_post import (
-    _parse_response as parse_create_response,
-)
-from ._generated_api_client.client import get_client
-from ._generated_api_client.models import (
-    ColumnType,
-    CreateDatasourceFromContentRequest,
-    DatasourceMetadata,
-)
 from ._utils.common import CreateMode, DropMode
 from ._utils.data_parsing import hf_dataset_from_torch
 from ._utils.tqdm_file_reader import TqdmFileReader
+from .client import DatasourceMetadata, orca_api
 
 
 def _upload_files_to_datasource(
@@ -55,8 +42,7 @@ def _upload_files_to_datasource(
     Returns:
         Metadata for the created datasource
     """
-    client = get_client()
-    files = []
+    files: list[tuple[Literal["files"], FileTypes]] = []
 
     # Calculate total size for all files
     total_size = sum(file_path.stat().st_size for file_path in file_paths)
@@ -65,16 +51,13 @@ def _upload_files_to_datasource(
         for file_path in file_paths:
             buffered_reader = open(file_path, "rb")
             tqdm_reader = TqdmFileReader(buffered_reader, pbar)
-            files.append(("files", (file_path.name, tqdm_reader)))
+            files.append(("files", (file_path.name, cast(bytes, tqdm_reader))))
 
         # Use manual HTTP request for file uploads
-        metadata = parse_create_response(
-            response=client.get_httpx_client().request(
-                method="post",
-                url="/datasource/upload",
-                files=files,
-                data={"name": name, "description": description},
-            )
+        metadata = orca_api.POST(
+            "/datasource/upload",
+            files=files,
+            data={"name": name, "description": description},
         )
 
     return metadata
@@ -129,19 +112,19 @@ class Datasource:
 
     def __init__(self, metadata: DatasourceMetadata):
         # for internal use only, do not document
-        self.id = metadata.id
-        self.name = metadata.name
-        self.length = metadata.length
-        self.created_at = metadata.created_at
-        self.updated_at = metadata.updated_at
-        self.description = metadata.description
+        self.id = metadata["id"]
+        self.name = metadata["name"]
+        self.length = metadata["length"]
+        self.created_at = datetime.fromisoformat(metadata["created_at"])
+        self.updated_at = datetime.fromisoformat(metadata["updated_at"])
+        self.description = metadata["description"]
         self.columns = {
-            column.name: (
-                f"enum({', '.join(f'{option!r}' for option in column.enum_options) if column.enum_options else ''})"
-                if column.type == ColumnType.ENUM
-                else "str" if column.type == ColumnType.STRING else column.type.value.lower()
+            column["name"]: (
+                f"enum({', '.join(f'{option!r}' for option in column['enum_options'] or []) if 'enum_options' in column else ''})"
+                if column["type"] == "ENUM"
+                else "str" if column["type"] == "STRING" else column["type"].lower()
             )
-            for column in metadata.columns
+            for column in metadata["columns"]
         }
 
     def __eq__(self, other) -> bool:
@@ -167,27 +150,16 @@ class Datasource:
 
         Returns:
             None
-
-        Raises:
-            RuntimeError: If the download fails.
         """
-
-        output_path = Path(output_path)
-        client = get_client().get_httpx_client()
-        url = f"/datasource/{self.id}/download"
-        response = client.get(url)
-        if response.status_code == 404:
-            raise LookupError(f"Datasource {self.id} not found")
-        if response.status_code != 200:
-            raise RuntimeError(f"Failed to download datasource: {response.status_code} {response.text}")
-
+        # TODO: add progress bar to the download
+        response = orca_api.GET("/datasource/{name_or_id}/download", params={"name_or_id": self.id}, parse_as=None)
         with tempfile.NamedTemporaryFile(suffix=".zip") as tmp_zip:
-            tmp_zip.write(response.content)
+            tmp_zip.write(response)
             tmp_zip.flush()
             with zipfile.ZipFile(tmp_zip.name, "r") as zf:
-                output_path.mkdir(parents=True, exist_ok=True)
+                Path(output_path).mkdir(parents=True, exist_ok=True)
                 for file in zf.namelist():
-                    out_file = output_path / Path(file).name
+                    out_file = Path(output_path) / Path(file).name
                     with zf.open(file) as af:
                         out_file.write_bytes(af.read())
 
@@ -317,10 +289,10 @@ class Datasource:
         if existing is not None:
             return existing
 
-        # Use the generated API client function for content creation
-        body = CreateDatasourceFromContentRequest(name=name, description=description, content=data)
-
-        metadata = create_datasource_from_content(body=body)
+        metadata = orca_api.POST(
+            "/datasource",
+            json={"name": name, "description": description, "content": data},
+        )
         return cls(metadata=metadata)
 
     @classmethod
@@ -351,10 +323,10 @@ class Datasource:
         if existing is not None:
             return existing
 
-        # Use the generated API client function for content creation
-        body = CreateDatasourceFromContentRequest(name=name, description=description, content=data)
-
-        metadata = create_datasource_from_content(body=body)
+        metadata = orca_api.POST(
+            "/datasource",
+            json={"name": name, "description": description, "content": data},
+        )
         return cls(metadata=metadata)
 
     @classmethod
@@ -410,16 +382,10 @@ class Datasource:
         parquet.write_table(pyarrow_table, buffer)
         parquet_bytes = buffer.getvalue()
 
-        client = get_client()
-
-        # Use manual HTTP request for file uploads
-        metadata = parse_create_response(
-            response=client.get_httpx_client().request(
-                method="post",
-                url="/datasource/upload",
-                files=[("files", ("data.parquet", parquet_bytes))],
-                data={"name": name, "description": description},
-            )
+        metadata = orca_api.POST(
+            "/datasource/upload",
+            files=[("files", ("data.parquet", parquet_bytes))],
+            data={"name": name, "description": description},
         )
 
         return cls(metadata=metadata)
@@ -471,12 +437,12 @@ class Datasource:
         return cls(metadata=metadata)
 
     @classmethod
-    def open(cls, name: str) -> Datasource:
+    def open(cls, name_or_id: str) -> Datasource:
         """
         Get a handle to a datasource by name or id in the OrcaCloud
 
         Params:
-            name: The name or unique identifier of the datasource to get
+            name_or_id: The name or unique identifier of the datasource to get
 
         Returns:
             A handle to the existing datasource in the OrcaCloud
@@ -484,7 +450,7 @@ class Datasource:
         Raises:
             LookupError: If the datasource does not exist
         """
-        return cls(get_datasource(name))
+        return cls(orca_api.GET("/datasource/{name_or_id}", params={"name_or_id": name_or_id}))
 
     @classmethod
     def exists(cls, name_or_id: str) -> bool:
@@ -511,7 +477,7 @@ class Datasource:
         Returns:
             A list of all datasource handles in the OrcaCloud
         """
-        return [cls(metadata) for metadata in list_datasources()]
+        return [cls(metadata) for metadata in orca_api.GET("/datasource")]
 
     @classmethod
     def drop(cls, name_or_id: str, if_not_exists: DropMode = "error") -> None:
@@ -527,7 +493,7 @@ class Datasource:
             LookupError: If the datasource does not exist and if_not_exists is `"error"`
         """
         try:
-            delete_datasource(name_or_id)
+            orca_api.DELETE("/datasource/{name_or_id}", params={"name_or_id": name_or_id})
             logging.info(f"Deleted datasource {name_or_id}")
         except LookupError:
             if if_not_exists == "error":