update client (#23)

* update client

* fixes

* fix content types

kscale/cli.py

@@ -3,15 +3,15 @@
 import click
 
 from kscale.utils.cli import recursive_help
-from kscale.web.kernel_images import cli as kernel_images_cli
+from kscale.web.kernels import cli as kernel_images_cli
 from kscale.web.krec import cli as krec_cli
 from kscale.web.pybullet import cli as pybullet_cli
 from kscale.web.urdf import cli as urdf_cli
 
 
 @click.group()
 def cli() -> None:
-    """Command line interface for interacting with the K-Scale store."""
+    """Command line interface for interacting with the K-Scale web API."""
     pass
 
 

kscale/conf.py

@@ -16,14 +16,14 @@ def get_path() -> Path:
 
 
 @dataclass
-class StoreSettings:
+class WWWSettings:
     api_key: str | None = field(default=None)
     cache_dir: str = field(default=II("oc.env:KSCALE_CACHE_DIR,'~/.kscale/cache/'"))
 
 
 @dataclass
 class Settings:
-    store: StoreSettings = field(default_factory=StoreSettings)
+    www: WWWSettings = field(default_factory=WWWSettings)
 
     def save(self) -> None:
         (dir_path := get_path()).mkdir(parents=True, exist_ok=True)

kscale/requirements.txt

@@ -11,3 +11,4 @@ pydantic
 
 # CLI
 click
+aiofiles

kscale/web/api.py

@@ -1,4 +1,4 @@
-"""Defines a common interface for the K-Scale Store API."""
+"""Defines a common interface for the K-Scale WWW API."""
 
 import asyncio
 from pathlib import Path

kscale/web/kernel_images.py → kscale/web/kernels.py

@@ -1,33 +1,33 @@
-"""Utility functions for managing kernel images in the K-Scale store."""
+"""Utility functions for managing kernel images in K-Scale WWW."""
 
 import hashlib
 import logging
 import shutil
 from pathlib import Path
 
+import aiofiles
 import click
 import httpx
 
 from kscale.utils.checksum import FileChecksum
 from kscale.utils.cli import coro
 from kscale.web.gen.api import SingleArtifactResponse
-from kscale.web.utils import get_api_key, get_artifact_dir, get_cache_dir
-from kscale.web.www_client import KScaleStoreClient
+from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_key, get_artifact_dir, get_cache_dir
+from kscale.web.www_client import KScaleWWWClient
 
 httpx_logger = logging.getLogger("httpx")
 httpx_logger.setLevel(logging.WARNING)
 
 logger = logging.getLogger(__name__)
 
 ALLOWED_SUFFIXES = {".img"}
-DEFAULT_UPLOAD_TIMEOUT = 300.0
 
 
 async def fetch_kernel_image_info(artifact_id: str, cache_dir: Path) -> SingleArtifactResponse:
     response_path = cache_dir / "response.json"
     if response_path.exists():
         return SingleArtifactResponse.model_validate_json(response_path.read_text())
-    async with KScaleStoreClient() as client:
+    async with KScaleWWWClient() as client:
         response = await client.get_artifact_info(artifact_id)
     response_path.write_text(response.model_dump_json())
     return response
@@ -46,7 +46,10 @@ async def download_kernel_image(artifact_id: str) -> Path:
         else:
             filename = cache_dir / original_name
 
-        headers = {"Authorization": f"Bearer {get_api_key()}", "Accept": "application/octet-stream"}
+        headers = {
+            "Authorization": f"Bearer {get_api_key()}",
+            "Accept": "application/octet-stream",
+        }
 
         if not filename.exists():
             logger.info("Downloading kernel image...")
@@ -56,10 +59,10 @@ async def download_kernel_image(artifact_id: str) -> Path:
                 async with client.stream("GET", artifact_url, headers=headers) as response:
                     response.raise_for_status()
 
-                    with open(filename, "wb") as f:
+                    async with aiofiles.open(filename, "wb") as f:
                         async for chunk in response.aiter_bytes():
                             FileChecksum.update_hash(sha256_hash, chunk)
-                            f.write(chunk)
+                            await f.write(chunk)
 
             logger.info("Kernel image downloaded to %s", filename)
         else:
@@ -107,7 +110,9 @@ async def remove_local_kernel_image(artifact_id: str) -> None:
 
 
 async def upload_kernel_image(
-    listing_id: str, image_path: Path, upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT
+    listing_id: str,
+    image_path: Path,
+    upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT,
 ) -> SingleArtifactResponse:
     """Upload a kernel image."""
     if image_path.suffix.lower() not in ALLOWED_SUFFIXES:
@@ -121,25 +126,27 @@ async def upload_kernel_image(
     logger.info("File name: %s", image_path.name)
     logger.info("File size: %.1f MB", file_size / 1024 / 1024)
 
-    async with KScaleStoreClient(upload_timeout=upload_timeout) as client:
+    async with KScaleWWWClient(upload_timeout=upload_timeout) as client:
         presigned_data = await client.get_presigned_url(
             listing_id=listing_id,
             file_name=image_path.name,
             checksum=checksum,
         )
 
         logger.info("Starting upload...")
-
-        with open(image_path, "rb") as f:
-            content = f.read()
-            async with httpx.AsyncClient() as http_client:
-                response = await http_client.put(
-                    presigned_data["upload_url"],
-                    content=content,
-                    headers={"Content-Type": "application/x-raw-disk-image"},
-                    timeout=upload_timeout,
-                )
-                response.raise_for_status()
+        async with httpx.AsyncClient() as http_client:
+            logger.info("Reading file content into memory...")
+            async with aiofiles.open(image_path, "rb") as f:
+                contents = await f.read()
+
+            logger.info("Uploading file content to %s", presigned_data["upload_url"])
+            response = await http_client.put(
+                presigned_data["upload_url"],
+                content=contents,
+                headers={"Content-Type": "application/x-raw-disk-image"},
+                timeout=upload_timeout,
+            )
+            response.raise_for_status()
 
         artifact_response: SingleArtifactResponse = await client.get_artifact_info(presigned_data["artifact_id"])
         logger.info("Uploaded artifact: %s", artifact_response.artifact_id)
@@ -156,7 +163,7 @@ async def upload_kernel_image_cli(
 
 @click.group()
 def cli() -> None:
-    """K-Scale Kernel Image Store CLI tool."""
+    """K-Scale Kernel Image CLI tool."""
     pass
 
 

kscale/web/krec.py

@@ -1,22 +1,29 @@
-"""Utility functions for managing K-Recs in the K-Scale store."""
+"""Utility functions for managing K-Recs in K-Scale WWW."""
 
 import asyncio
 import json
 import logging
 from pathlib import Path
 
+import aiofiles
 import click
 import httpx
 
 from kscale.utils.cli import coro
 from kscale.web.gen.api import UploadKRecRequest
-from kscale.web.utils import get_api_key, get_artifact_dir
-from kscale.web.www_client import KScaleStoreClient
+from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_key, get_artifact_dir
+from kscale.web.www_client import KScaleWWWClient
 
 logger = logging.getLogger(__name__)
 
 
-async def upload_krec(robot_id: str, file_path: Path, name: str, description: str | None = None) -> str:
+async def upload_krec(
+    robot_id: str,
+    file_path: Path,
+    name: str,
+    description: str | None = None,
+    upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT,
+) -> str:
     if not file_path.exists():
         raise FileNotFoundError(f"File not found: {file_path}")
 
@@ -25,7 +32,7 @@ async def upload_krec(robot_id: str, file_path: Path, name: str, description: st
     logger.info("File name: %s", file_path.name)
     logger.info("File size: %.1f MB", file_size / 1024 / 1024)
 
-    async with KScaleStoreClient() as client:
+    async with KScaleWWWClient(upload_timeout=upload_timeout) as client:
         create_response = await client.create_krec(
             UploadKRecRequest(
                 robot_id=robot_id,
@@ -36,16 +43,19 @@ async def upload_krec(robot_id: str, file_path: Path, name: str, description: st
 
         logger.info("Initialized K-Rec upload with ID: %s", create_response["krec_id"])
         logger.info("Starting upload...")
-
-        with open(file_path, "rb") as f:
-            content = f.read()
-            async with httpx.AsyncClient() as http_client:
-                response = await http_client.put(
-                    create_response["upload_url"],
-                    content=content,
-                    headers={"Content-Type": "application/octet-stream"},
-                )
-                response.raise_for_status()
+        async with httpx.AsyncClient() as http_client:
+            logger.info("Reading file content into memory...")
+            async with aiofiles.open(file_path, "rb") as f:
+                contents = await f.read()
+
+            logger.info("Uploading file content to %s", create_response["upload_url"])
+            response = await http_client.put(
+                create_response["upload_url"],
+                content=contents,
+                headers={"Content-Type": "video/x-matroska"},
+                timeout=upload_timeout,
+            )
+            response.raise_for_status()
 
         logger.info("Successfully uploaded K-Rec: %s", create_response["krec_id"])
         return create_response["krec_id"]
@@ -61,7 +71,7 @@ async def fetch_krec_info(krec_id: str, cache_dir: Path) -> dict:
     if response_path.exists():
         return json.loads(response_path.read_text())
 
-    async with KScaleStoreClient() as client:
+    async with KScaleWWWClient() as client:
         try:
             response = await client.get_krec_info(krec_id)
 

kscale/web/urdf.py

@@ -1,4 +1,4 @@
-"""Utility functions for managing artifacts in the K-Scale store."""
+"""Utility functions for managing artifacts in K-Scale WWW."""
 
 import logging
 import shutil
@@ -12,7 +12,7 @@
 from kscale.utils.cli import coro
 from kscale.web.gen.api import SingleArtifactResponse, UploadArtifactResponse
 from kscale.web.utils import get_api_key, get_artifact_dir, get_cache_dir
-from kscale.web.www_client import KScaleStoreClient
+from kscale.web.www_client import KScaleWWWClient
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -34,7 +34,7 @@ async def fetch_urdf_info(artifact_id: str, cache_dir: Path) -> SingleArtifactRe
     response_path = cache_dir / "response.json"
     if response_path.exists():
         return SingleArtifactResponse.model_validate_json(response_path.read_text())
-    async with KScaleStoreClient() as client:
+    async with KScaleWWWClient() as client:
         response = await client.get_artifact_info(artifact_id)
     response_path.write_text(response.model_dump_json())
     return response
@@ -129,7 +129,7 @@ async def remove_local_urdf(artifact_id: str) -> None:
 async def upload_urdf(listing_id: str, root_dir: Path) -> UploadArtifactResponse:
     tarball_path = create_tarball(root_dir, "robot.tgz", get_artifact_dir(listing_id))
 
-    async with KScaleStoreClient() as client:
+    async with KScaleWWWClient() as client:
         response = await client.upload_artifact(listing_id, str(tarball_path))
 
     logger.info("Uploaded artifacts: %s", [artifact.artifact_id for artifact in response.artifacts])

kscale/web/utils.py

@@ -1,29 +1,31 @@
-"""Utility functions for interacting with the K-Scale Store API."""
+"""Utility functions for interacting with the K-Scale WWW API."""
 
 import os
 from pathlib import Path
 
 from kscale.conf import Settings
 
+DEFAULT_UPLOAD_TIMEOUT = 300.0  # 5 minutes
+
 
 def get_api_root() -> str:
-    """Returns the base URL for the K-Scale Store API.
+    """Returns the base URL for the K-Scale WWW API.
 
     This can be overridden when targetting a different server.
 
     Returns:
-        The base URL for the K-Scale Store API.
+        The base URL for the K-Scale WWW API.
     """
     return os.getenv("KSCALE_API_ROOT", "https://api.kscale.dev")
 
 
 def get_api_key() -> str:
-    """Returns the API key for the K-Scale Store API.
+    """Returns the API key for the K-Scale WWW API.
 
     Returns:
-        The API key for the K-Scale Store API.
+        The API key for the K-Scale WWW API.
     """
-    api_key = Settings.load().store.api_key
+    api_key = Settings.load().www.api_key
     if api_key is None:
         api_key = os.getenv("KSCALE_API_KEY")
     if not api_key:
@@ -36,7 +38,7 @@ def get_api_key() -> str:
 
 def get_cache_dir() -> Path:
     """Returns the cache directory for artifacts."""
-    return Path(Settings.load().store.cache_dir).expanduser().resolve()
+    return Path(Settings.load().www.cache_dir).expanduser().resolve()
 
 
 def get_artifact_dir(artifact_id: str) -> Path:

kscale/web/www_client.py

@@ -1,4 +1,4 @@
-"""Defines a typed client for the K-Scale Store API."""
+"""Defines a typed client for the K-Scale WWW API."""
 
 import logging
 from pathlib import Path
@@ -16,13 +16,13 @@
     UploadArtifactResponse,
     UploadKRecRequest,
 )
-from kscale.web.utils import get_api_key, get_api_root
+from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_key, get_api_root
 
 logger = logging.getLogger(__name__)
 
 
-class KScaleStoreClient:
-    def __init__(self, base_url: str = get_api_root(), upload_timeout: float = 300.0) -> None:
+class KScaleWWWClient:
+    def __init__(self, base_url: str = get_api_root(), upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT) -> None:
         self.base_url = base_url
         self.upload_timeout = upload_timeout
         self._client: httpx.AsyncClient | None = None
@@ -89,7 +89,7 @@ async def close(self) -> None:
             await self._client.aclose()
             self._client = None
 
-    async def __aenter__(self) -> "KScaleStoreClient":
+    async def __aenter__(self) -> "KScaleWWWClient":
         return self
 
     async def __aexit__(