Merge pull request #8 from kscalelabs/api-improvements

api improvements

README.md

@@ -30,3 +30,32 @@ This is a command line tool for interacting with various services provided by K-
 ```bash
 pip install kscale
 ```
+
+## Usage
+
+### CLI
+
+Download a URDF from the K-Scale Store:
+
+```bash
+kscale urdf download <artifact_id>
+```
+
+Upload a URDF to the K-Scale Store:
+
+```bash
+kscale urdf upload <artifact_id> <root_dir>
+```
+
+### Python API
+
+Reference a URDF by ID from the K-Scale Store:
+
+```python
+from kscale import KScale
+
+async def main():
+  kscale = KScale()
+  urdf_dir_path = await kscale.store.urdf("123456")
+  print(urdf_dir_path)
+```

kscale/__init__.py

@@ -1 +1,5 @@
+"""Defines the common interface for the K-Scale Python API."""
+
 __version__ = "0.0.4"
+
+from kscale.api import KScale

kscale/api.py

@@ -0,0 +1,14 @@
+"""Defines common functionality for the K-Scale API."""
+
+from kscale.store.api import StoreAPI
+from kscale.utils.api_base import APIBase
+
+
+class KScale(
+    StoreAPI,
+    APIBase,
+):
+    """Defines a common interface for the K-Scale API."""
+
+    def __init__(self, api_key: str | None = None) -> None:
+        self.api_key = api_key

kscale/conf.py

@@ -17,7 +17,7 @@ def get_path() -> Path:
 
 @dataclass
 class StoreSettings:
-    api_key: str = field(default=II("oc.env:KSCALE_API_KEY,"))
+    api_key: str | None = field(default=None)
     cache_dir: str = field(default=II("oc.env:KSCALE_CACHE_DIR,'~/.kscale/cache/'"))
 
 

kscale/store/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.0.1"

kscale/store/api.py

@@ -0,0 +1,20 @@
+"""Defines a common interface for the K-Scale Store API."""
+
+from pathlib import Path
+
+from kscale.store.urdf import download_urdf
+from kscale.utils.api_base import APIBase
+
+
+class StoreAPI(APIBase):
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+    ) -> None:
+        super().__init__()
+
+        self.api_key = api_key
+
+    async def urdf(self, artifact_id: str) -> Path:
+        return await download_urdf(artifact_id)

kscale/store/client.py

@@ -0,0 +1,79 @@
+"""Defines a typed client for the K-Scale Store API."""
+
+import logging
+from types import TracebackType
+from typing import Any, Dict, Type
+from urllib.parse import urljoin
+
+import httpx
+from pydantic import BaseModel
+
+from kscale.store.gen.api import (
+    NewListingRequest,
+    NewListingResponse,
+    SingleArtifactResponse,
+    UploadArtifactResponse,
+)
+from kscale.store.utils import API_ROOT, get_api_key
+
+logger = logging.getLogger(__name__)
+
+
+class KScaleStoreClient:
+    def __init__(self, base_url: str = API_ROOT) -> None:
+        self.base_url = base_url
+        self.client = httpx.AsyncClient(
+            base_url=self.base_url,
+            headers={"Authorization": f"Bearer {get_api_key()}"},
+        )
+
+    async def _request(
+        self,
+        method: str,
+        endpoint: str,
+        *,
+        params: Dict[str, Any] | None = None,
+        data: BaseModel | None = None,
+        files: Dict[str, Any] | None = None,
+    ) -> Dict[str, Any]:
+        url = urljoin(self.base_url, endpoint)
+        kwargs: Dict[str, Any] = {"params": params}
+
+        if data:
+            kwargs["json"] = data.dict(exclude_unset=True)
+        if files:
+            kwargs["files"] = files
+
+        response = await self.client.request(method, url, **kwargs)
+        if response.is_error:
+            logger.error(f"Error response from K-Scale Store: {response.text}")
+        response.raise_for_status()
+        return response.json()
+
+    async def get_artifact_info(self, artifact_id: str) -> SingleArtifactResponse:
+        data = await self._request("GET", f"/artifacts/info/{artifact_id}")
+        return SingleArtifactResponse(**data)
+
+    async def upload_artifact(self, listing_id: str, file_path: str) -> UploadArtifactResponse:
+        with open(file_path, "rb") as f:
+            files = {"files": (f.name, f, "application/gzip")}
+            data = await self._request("POST", f"/artifacts/upload/{listing_id}", files=files)
+        return UploadArtifactResponse(**data)
+
+    async def create_listing(self, request: NewListingRequest) -> NewListingResponse:
+        data = await self._request("POST", "/listings", data=request)
+        return NewListingResponse(**data)
+
+    async def close(self) -> None:
+        await self.client.aclose()
+
+    async def __aenter__(self) -> "KScaleStoreClient":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.close()

kscale/store/gen/api.py

@@ -2,12 +2,12 @@
 
 # generated by datamodel-codegen:
 #   filename:  openapi.json
-#   timestamp: 2024-08-30T02:21:23+00:00
+#   timestamp: 2024-09-04T04:33:58+00:00
 
 from __future__ import annotations
 
 from enum import Enum
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from pydantic import BaseModel, EmailStr, Field
 
@@ -21,8 +21,9 @@ class AuthResponse(BaseModel):
     api_key: str = Field(..., title="Api Key")
 
 
-class BodySetUrdfUrdfUploadListingIdPost(BaseModel):
-    file: bytes = Field(..., title="File")
+class BodyPullOnshapeDocumentOnshapePullListingIdGet(BaseModel):
+    suffix_to_joint_effort: Optional[Dict[str, float]] = Field(None, title="Suffix To Joint Effort")
+    suffix_to_joint_velocity: Optional[Dict[str, float]] = Field(None, title="Suffix To Joint Velocity")
 
 
 class BodyUploadArtifactsUploadListingIdPost(BaseModel):
@@ -57,6 +58,7 @@ class GetListingResponse(BaseModel):
     views: int = Field(..., title="Views")
     score: int = Field(..., title="Score")
     user_vote: Optional[bool] = Field(..., title="User Vote")
+    creator_id: str = Field(..., title="Creator Id")
 
 
 class GetTokenResponse(BaseModel):
@@ -254,16 +256,6 @@ class UploadArtifactResponse(BaseModel):
     artifacts: List[SingleArtifactResponse] = Field(..., title="Artifacts")
 
 
-class UrdfInfo(BaseModel):
-    artifact_id: str = Field(..., title="Artifact Id")
-    url: str = Field(..., title="Url")
-
-
-class UrdfResponse(BaseModel):
-    urdf: Optional[UrdfInfo]
-    listing_id: str = Field(..., title="Listing Id")
-
-
 class UserInfoResponseItem(BaseModel):
     id: str = Field(..., title="Id")
     email: str = Field(..., title="Email")