Bug fixes and enhancements (#4)

Author: jiahuei

.github/workflows/ci.yml

@@ -28,20 +28,10 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
-      #- name: Test API Key
-      #  run: |
-      #    response=$(curl --fail -s -H "Authorization: Bearer ${OPENAI_API_KEY}" https://api.openai.com/v1/engines)
-      #    echo "$response"
-      #    if [ $? -ne 0 ]; then
-      #      echo "API key test failed!"
-      #      exit 1
-      #    fi
-      #  env:
-      #    OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-
-      - name: Setup
+      - name: Remove cloud-only modules and install Python client
         run: |
           set -e
+          bash scripts/remove_cloud_modules.sh
           cd clients/python
           python -m pip install .[all]
 
@@ -98,20 +88,6 @@ jobs:
           COMPOSE_DOCKER_CLI_BUILD: 1
           DOCKER_BUILDKIT: 1
 
-      #- name: Test list models
-      #  run: |
-      #    response=$(curl --fail -s http://localhost:6969/api/v1/models)
-      #    echo "$response"
-      #    if [ $? -ne 0 ]; then
-      #      echo "list models test failed!"
-      #      exit 1
-      #    fi
-
-      #- name: Get docker compose logs
-      #  if: always()
-      #  run: |
-      #    docker compose -f docker/compose.cpu.yml logs
-
       - name: Pytest
         run: |
           set -e

CHANGELOG.md

@@ -2,12 +2,10 @@
 
 All notable changes to this project will be documented in this file.
 
-The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). [An example](https://github.com/standard/standard/blob/master/CHANGELOG.md).
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
-## [Unreleased]
+The version number mentioned here refers to the cloud version.
 
-This is for a future release.
+## [v0.1] - 2024-06-03
 
-### CHANGED
-
-### FIXED
+This is our first release 🚀

README.md

@@ -8,7 +8,7 @@
 
 <!-- prettier-ignore -->
 > [!TIP]
-> [Explore our docs](https://docs.jamaibase.com)
+> [Explore our docs](#explore-the-documentation)
 
 ## Overview
 
@@ -141,10 +141,10 @@ Get free LLM tokens on JamAI Base Cloud. [Sign up now.](https://cloud.jamaibase.
    $ docker compose -f docker/compose.nvidia.yml up --quiet-pull -d
    ```
 
-   <!-- prettier-ignore -->
-   > [!TIP]
-   > By default, frontend and backend are accessible at ports 4000 and 6969.
-   > You can change the ports exposed to host by setting env var like so `API_PORT=6970 FRONTEND_PORT=4001 docker compose -f docker/compose.cpu.yml up --quiet-pull -d`
+<!-- prettier-ignore -->
+> [!TIP]
+> By default, frontend and backend are accessible at ports 4000 and 6969.
+> You can change the ports exposed to host by setting env var like so `API_PORT=6970 FRONTEND_PORT=4001 docker compose -f docker/compose.cpu.yml up --quiet-pull -d`
 
 4. Try the command below in your terminal, or open your browser and go to `localhost:4000`.
 
@@ -156,6 +156,8 @@ Get free LLM tokens on JamAI Base Cloud. [Sign up now.](https://cloud.jamaibase.
 
 - [API Documentation](https://jamaibase.readme.io)
 - [Platform Documentation](https://docs.jamaibase.com)
+- [Changelog](CHANGELOG.md)
+- [Versioning](VERSIONING.md)
 
 ## Examples
 

VERSIONING.md

@@ -0,0 +1,20 @@
+# Versioning
+
+This document describes the versioning policy for this repository.
+
+# TLDR
+
+We follow [Semantic Versioning](https://semver.org/).
+
+# Current Client SDK Versions (v0.x)
+
+We will start our versioning at `v0.x.x`.
+
+Major version zero (0.y.z) is for initial development. Anything may change at any time. The public API should not be considered stable.
+
+# Future Client SDK Versions (v1.x)
+
+SDK versions will follow this structure `<major>.<minor>.<patch>`:
+
+- `major` is incremented for any backwards-incompatible updates made to our [Cloud service / backend](https://cloud.jamaibase.com/) and hence our SDK.
+- `minor` or `patch` is incremented for any backwards-compatible feature update, bug fix, enhancements made to the SDK libraries.

clients/python/src/jamaibase/protocol.py

@@ -1206,7 +1206,13 @@ def handle_nulls_and_validate(self) -> Self:
                 if k in failed_cols:
                     d[k], state["original"] = None, d[k]
                 if d[k] is None:
-                    if col.dtype == DtypeEnum.str_:
+                    if col.dtype == DtypeEnum.int_:
+                        d[k] = 0
+                    elif col.dtype == DtypeEnum.float_:
+                        d[k] = 0.0
+                    elif col.dtype == DtypeEnum.bool_:
+                        d[k] = False
+                    elif col.dtype == DtypeEnum.str_:
                         # Store null string as ""
                         # https://github.com/lancedb/lancedb/issues/1160
                         d[k] = ""

clients/python/src/jamaibase/utils/io.py

@@ -14,9 +14,10 @@
 import orjson
 import srsly
 import toml
-from jamaibase.utils.types import JSONInput, JSONOutput
 from PIL import ExifTags, Image
 
+from jamaibase.utils.types import JSONInput, JSONOutput
+
 logger = logging.getLogger(__name__)
 
 
@@ -49,7 +50,7 @@ def dump_json(data: JSONInput, path: str, **kwargs) -> str:
     Args:
         data (JSONInput): The data.
         path (str): Path to the file.
-        **kwargs: Other keyword arguments to pass into `srsly.write_json`.
+        **kwargs: Other keyword arguments to pass into `orjson.dumps`.
 
     Returns:
         path (str): Path to the file.

clients/python/src/jamaibase/version.py

@@ -1 +1 @@
-__version__ = "0.0.1"
+__version__ = "0.1.0"

scripts/remove_cloud_modules.sh

@@ -1,9 +1,7 @@
 #!/usr/bin/env bash
 
-set -e
-
 find . -type f -name "cloud*.py" -delete
 find . -type f -name "compose.*.cloud.yml" -delete
 find . -type d -name "(cloud)" -exec rm -rf {} +
-rm services/app/ecosystem.config.cjs
-rm services/app/ecosystem.json
+rm -f services/app/ecosystem.config.cjs
+rm -f services/app/ecosystem.json

services/api/pyproject.toml

@@ -78,13 +78,14 @@ dependencies = [
     "lancedb~=0.6.13",
     "langchain-community~=0.0.25",
     "langchain~=0.1.10",
-    "litellm~=1.35.15",
+    "litellm~=1.40.0",
     "loguru~=0.7.2",
     "matplotlib",
     "numpy~=1.26.4",
-    "openai~=1.23.1",
+    "openai~=1.30.5",
     "openmeter~=1.0.0b50",
     "orjson~=3.9.15",
+    "Pillow~=10.3.0",
     "pyarrow==15.0.0",
     "pycryptodomex~=3.20.0",
     "pydantic-settings>=2.2.1",
@@ -93,9 +94,11 @@ dependencies = [
     "python-multipart~=0.0.6",
     "redis[hiredis]~=5.0.4",
     "sqlmodel~=0.0.16",
+    "srsly~=2.4.8",
     "stripe~=8.5.0",
     "tantivy~=0.21.0",
     "tiktoken~=0.6.0",
+    "toml~=0.10.2",
     "tqdm~=4.66.2",
     "typer[all]~=0.9.0",
     "typing_extensions>=4.10.0",

services/api/src/owl/db/gen_table.py

@@ -12,7 +12,6 @@
 from filelock import FileLock
 from lancedb.table import LanceTable
 from loguru import logger
-from pydantic import BaseModel
 from sqlalchemy import desc
 from sqlmodel import Session, SQLModel, select
 from typing_extensions import Self
@@ -67,7 +66,7 @@ def __init__(
         self.db_url = Path(db_url)
         self.vector_db_url = Path(vector_db_url)
 
-    def lock(self, name: str, timeout: int = 60):
+    def lock(self, name: str, timeout: int = 5):
         name = f"{self.lock_name_prefix}/{name}.lock"
         self.locks[name] = self.locks.get(name, FileLock(name, timeout=timeout))
         return self.locks[name]

services/api/src/owl/entrypoints/api.py

@@ -3,6 +3,7 @@
 
 ```shell
 $ python -m owl.entrypoints.api
+$ JAMAI_API_BASE=http://localhost:6969/api TZ=Asia/Singapore python -m owl.entrypoints.api
 ```
 """
 
@@ -16,6 +17,7 @@
 from fastapi.exceptions import RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import ORJSONResponse
+from filelock import Timeout
 from loguru import logger
 from pydantic import SecretStr
 from pydantic_settings import BaseSettings, SettingsConfigDict
@@ -226,8 +228,8 @@ async def authenticate(request: Request, call_next):
     # Defaults
     project_id, org_id, org_tier = CONFIG.default_project, CONFIG.default_org, "free"
     token, openmeter_id = "", "default"
-    quota_reset_at = ""
-    quotas = defaultdict(lambda: 1.0)
+    quotas, quota_reset_at = defaultdict(lambda: 1.0), ""
+    db_storage_gb = file_storage_gb = 0.0
 
     # --- OSS Mode --- #
     if CONFIG.service_key_plain == "":
@@ -380,9 +382,11 @@ async def authenticate(request: Request, call_next):
                     )
                 org_id, external_keys = org_info.id, org_info.external_keys
             openmeter_id = org_info.openmeter_id
-            quota_reset_at = org_info.quota_reset_at
             org_tier = org_info.tier
             quotas = org_info.quotas
+            quota_reset_at = org_info.quota_reset_at
+            db_storage_gb = org_info.db_storage_gb
+            file_storage_gb = org_info.file_storage_gb
             if openmeter_id is None:
                 logger.warning(
                     f"{request.state.id} - Organization {org_id} does not have OpenMeter ID."
@@ -399,14 +403,14 @@ async def authenticate(request: Request, call_next):
     # --- Set request state and headers --- #
     request.state.org_id = org_id
     request.state.project_id = project_id
-    request.state.api_key = token
-    request.state.openmeter_id = openmeter_id
     request.state.billing_manager = BillingManager(
         request=request,
+        openmeter_id=openmeter_id,
         quotas=quotas,
         quota_reset_at=quota_reset_at,
+        db_storage_gb=db_storage_gb,
+        file_storage_gb=file_storage_gb,
         organization_tier=org_tier,
-        openmeter_id=openmeter_id,
         organization_id=org_id,
         project_id=project_id,
         api_key=token,
@@ -457,6 +461,21 @@ async def health() -> Response:
 # --- Order of handlers does not matter --- #
 
 
+@app.exception_handler(Timeout)
+async def write_lock_timeout_exc_handler(request: Request, exc: Timeout):
+    return ORJSONResponse(
+        status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+        content={
+            "error": "write_lock_timeout",
+            "message": "This table is currently busy. Please try again later.",
+            "detail": str(exc),
+            "request_id": request.state.id,
+            "exception": exc.__class__.__name__,
+        },
+        headers={"Retry-After": 10},
+    )
+
+
 @app.exception_handler(UpgradeTierError)
 async def upgrade_tier_exc_handler(request: Request, exc: UpgradeTierError):
     return ORJSONResponse(

services/api/src/owl/routers/gen_table.py

@@ -1,7 +1,8 @@
-import asyncio
 import pathlib
+from asyncio import sleep
 from datetime import timedelta
 from hashlib import blake2b
+from time import perf_counter
 from typing import Annotated, Any
 
 import numpy as np
@@ -122,7 +123,7 @@ async def periodic_reindex():
     lock = FileLock(f"{config.owl_db_dir}/periodic_reindex.lock", blocking=False)
     try:
         with lock:
-            logger.info("Periodic Lance re-indexing started.")
+            t0 = perf_counter()
             num_ok = num_skipped = num_failed = 0
             for session, table, meta, table_path in _iter_all_tables():
                 if session is None:
@@ -136,14 +137,18 @@ async def periodic_reindex():
                 except Exception:
                     logger.exception(f"Periodic Lance re-indexing failed for table: {table_path}")
                     num_failed += 1
+            t = perf_counter() - t0
+            # Hold the lock for a while to block other workers
+            await sleep(max(0.0, (config.owl_reindex_period_sec - t) * 0.5))
         logger.info(
             (
-                "Periodic Lance re-indexing completed "
-                f"({num_ok:,d} OK, {num_skipped:,d} skipped, {num_failed:,d} failed)."
+                f"Periodic Lance re-indexing completed (t={t:,.3f} s, "
+                f"{num_ok:,d} OK, {num_skipped:,d} skipped, {num_failed:,d} failed)."
             )
         )
     except Timeout:
-        logger.info("Periodic Lance re-indexing skipped.")
+        # logger.info("Periodic Lance re-indexing skipped.")
+        pass
     except Exception:
         logger.exception("Periodic Lance re-indexing encountered an error.")
 
@@ -154,7 +159,7 @@ async def periodic_optimize():
     lock = FileLock(f"{config.owl_db_dir}/periodic_optimization.lock", blocking=False)
     try:
         with lock:
-            logger.info("Periodic Lance optimization started.")
+            t0 = perf_counter()
             num_ok = num_skipped = num_failed = 0
             for _, table, meta, table_path in _iter_all_tables():
                 done = True
@@ -177,14 +182,18 @@ async def periodic_optimize():
                 except Exception:
                     logger.exception(f"Periodic Lance optimization failed for table: {table_path}")
                     num_failed += 1
+            t = perf_counter() - t0
+            # Hold the lock for a while to block other workers
+            await sleep(max(0.0, (config.owl_reindex_period_sec - t) * 0.5))
         logger.info(
             (
-                "Periodic Lance optimization completed "
-                f"({num_ok:,d} OK, {num_skipped:,d} skipped, {num_failed:,d} failed)."
+                f"Periodic Lance optimization completed (t={t:,.3f} s, "
+                f"{num_ok:,d} OK, {num_skipped:,d} skipped, {num_failed:,d} failed)."
             )
         )
     except Timeout:
-        logger.info("Periodic Lance optimization skipped.")
+        # logger.info("Periodic Lance optimization skipped.")
+        pass
     except Exception:
         logger.exception("Periodic Lance optimization encountered an error.")
 

services/api/tests/test_cloud_client.py

@@ -81,19 +81,13 @@ def test_get_and_list_organizations(client_cls: Type[Owl]):
 @pytest.mark.parametrize("client_cls", CLIENT_CLS)
 def test_create_flow(client_cls: Type[Owl]):
     owl = client_cls()
-    # response = owl.delete_user(USER_ID_A)
-    # assert isinstance(response, OkResponse)
-    # response = owl.delete_user(USER_ID_B)
-    # assert isinstance(response, OkResponse)
+    response = owl.delete_user(USER_ID_A)
+    assert isinstance(response, OkResponse)
+    response = owl.delete_user(USER_ID_B)
+    assert isinstance(response, OkResponse)
     # response = owl.delete_organization("org_f12d46652cfa3120006e44e2")
     # assert isinstance(response, OkResponse)
 
-    orgs = owl.list_organizations().items
-    for org in orgs:
-        org = owl.refresh_quota(org.id)
-        print(org)
-    exit()
-
     # Create user
     duncan = _create_user(owl)
     print(f"User created: {duncan}\n")

services/app/build.sh

@@ -22,9 +22,7 @@ vite build
 
 # Copy the build files to the build directory
 rm -rf build
-mkdir build
-cp -r temp/* build/
-rm -rf temp
+mv temp build
 
 # Reload PM2 app if not in dev mode
 if [ $DEV_MODE -eq 0 ]; then