orca-sdk 0.0.97__py3-none-any.whl → 0.0.100__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

@@ -1,40 +1,28 @@
 from __future__ import annotations
 
 import logging
+import os
 import tempfile
 import zipfile
 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 +43,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 +52,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 +113,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:
@@ -158,39 +142,6 @@ class Datasource:
             + "})"
         )
 
-    def download(self, output_path: str | PathLike) -> None:
-        """
-        Download the datasource as a ZIP and extract them to a specified path.
-
-        Params:
-            output_path: The local file path or directory where the downloaded files will be saved.
-
-        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}")
-
-        with tempfile.NamedTemporaryFile(suffix=".zip") as tmp_zip:
-            tmp_zip.write(response.content)
-            tmp_zip.flush()
-            with zipfile.ZipFile(tmp_zip.name, "r") as zf:
-                output_path.mkdir(parents=True, exist_ok=True)
-                for file in zf.namelist():
-                    out_file = output_path / Path(file).name
-                    with zf.open(file) as af:
-                        out_file.write_bytes(af.read())
-
     @classmethod
     def from_hf_dataset(
         cls, name: str, dataset: Dataset, if_exists: CreateMode = "error", description: str | None = None
@@ -317,10 +268,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 +302,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 +361,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 +416,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 +429,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 +456,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 +472,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":
@@ -535,3 +480,45 @@ class Datasource:
 
     def __len__(self) -> int:
         return self.length
+
+    def download(
+        self, output_dir: str | PathLike, file_type: Literal["hf_dataset", "json", "csv"] = "hf_dataset"
+    ) -> None:
+        """
+        Download the datasource to a specified path in the specified format type
+
+        Params:
+            output_dir: The local directory where the downloaded file will be saved.
+            file_type: The type of file to download.
+
+        Returns:
+            None
+        """
+        extension = "zip" if file_type == "hf_dataset" else file_type
+        output_path = Path(output_dir) / f"{self.name}.{extension}"
+        with open(output_path, "wb") as download_file:
+            with orca_api.stream("GET", f"/datasource/{self.id}/download", params={"file_type": file_type}) as response:
+                total_chunks = int(response.headers["X-Total-Chunks"]) if "X-Total-Chunks" in response.headers else None
+                with tqdm(desc=f"Downloading", total=total_chunks, disable=total_chunks is None) as progress:
+                    for chunk in response.iter_bytes():
+                        download_file.write(chunk)
+                        progress.update(1)
+
+        # extract the zip file
+        if extension == "zip":
+            extract_dir = Path(output_dir) / self.name
+            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}")
+        else:
+            logging.info(f"Downloaded {output_path}")
+
+    def to_list(self) -> list[dict]:
+        """
+        Convert the datasource to a list of dictionaries.
+
+        Returns:
+            A list of dictionaries representation of the datasource.
+        """
+        return orca_api.GET("/datasource/{name_or_id}/download", params={"name_or_id": self.id, "file_type": "json"})

orca_sdk/datasource_test.py

@@ -1,10 +1,14 @@
+import json
 import os
 import tempfile
+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
 
 from .datasource import Datasource
 
@@ -99,13 +103,6 @@ def test_drop_datasource_invalid_input():
         Datasource.drop("not valid id")
 
 
-def test_download_datasource(datasource):
-    with tempfile.TemporaryDirectory() as temp_dir:
-        output_path = os.path.join(temp_dir, "datasource.zip")
-        datasource.download(output_path)
-        assert os.path.exists(output_path)
-
-
 def test_from_list():
     # Test creating datasource from list of dictionaries
     data = [
@@ -296,3 +293,40 @@ def test_from_disk_already_exists():
             assert datasource2.name == datasource1.name
         finally:
             os.unlink(f.name)
+
+
+def test_to_list(hf_dataset, datasource):
+    assert datasource.to_list() == hf_dataset.to_list()
+
+
+def test_download_datasource(hf_dataset, datasource):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Dataset download
+        datasource.download(temp_dir)
+        downloaded_hf_dataset_dir = f"{temp_dir}/{datasource.name}"
+        assert os.path.exists(downloaded_hf_dataset_dir)
+        assert os.path.isdir(downloaded_hf_dataset_dir)
+        assert not os.path.exists(f"{downloaded_hf_dataset_dir}.zip")
+        dataset_from_downloaded_hf_dataset = Dataset.load_from_disk(downloaded_hf_dataset_dir)
+        assert dataset_from_downloaded_hf_dataset.column_names == hf_dataset.column_names
+        assert dataset_from_downloaded_hf_dataset.to_dict() == hf_dataset.to_dict()
+
+        # JSON download
+        datasource.download(temp_dir, file_type="json")
+        downloaded_json_file = f"{temp_dir}/{datasource.name}.json"
+        assert os.path.exists(downloaded_json_file)
+        with open(downloaded_json_file, "r") as f:
+            content = json.load(f)
+            assert content == hf_dataset.to_list()
+
+        # CSV download
+        datasource.download(temp_dir, file_type="csv")
+        downloaded_csv_file = f"{temp_dir}/{datasource.name}.csv"
+        assert os.path.exists(downloaded_csv_file)
+        dataset_from_downloaded_csv = cast(Dataset, Dataset.from_csv(downloaded_csv_file))
+        assert dataset_from_downloaded_csv.column_names == hf_dataset.column_names
+        assert (
+            dataset_from_downloaded_csv.remove_columns("score").to_dict()
+            == hf_dataset.remove_columns("score").to_dict()
+        )
+        assert np.allclose(dataset_from_downloaded_csv["score"], hf_dataset["score"])