[Backend] Various updates ; Draft polylm library (#879)

Python SDK - jamaibase

- Support context manager

Backend - owl (API server)

- Log invitation code rejection reason
- Remove unsupported DB dialects
- Only allow DB init in dev
- Bump `owl` dependencies
- Move owl scripts
- Draft `polylm` library

---

Co-authored-by: AmirBalwel <amir.balwel@embeddedllm.com>

Author: jiahuei

clients/python/pyproject.toml

@@ -71,7 +71,7 @@ unfixable = ["B"]
 "**/{tests,docs,tools}/*" = ["E402"]
 
 [tool.ruff.lint.isort]
-known-first-party = ["jamaibase", "owl"]
+known-first-party = ["jamaibase", "owl", "polylm"]
 
 [tool.ruff.lint.flake8-bugbear]
 # Allow default arguments like, e.g., `data: List[str] = fastapi.Query(None)`.

clients/python/src/jamaibase/client.py

@@ -4,7 +4,7 @@
 from datetime import datetime
 from os.path import basename, split
 from time import perf_counter
-from typing import Any, AsyncGenerator, BinaryIO, Generator, Literal, Type
+from typing import Any, AsyncGenerator, BinaryIO, Generator, Literal, Self, Type
 from urllib.parse import quote
 from warnings import warn
 
@@ -192,6 +192,17 @@ async def close(self) -> None:
         """
         await self.http_client.aclose()
 
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+        if exc_type:
+            logger.error(f"Client exited with exception: {exc_type.__name__}: {exc_value}")
+        try:
+            await self.close()
+        except Exception as e:
+            logger.error(f"Client failed to close: {e}")
+
     @staticmethod
     def _filter_params(params: dict[str, Any] | BaseModel | None) -> dict[str, Any] | None:
         """

clients/python/src/jamaibase/types/lm.py

@@ -6,7 +6,6 @@
     AfterValidator,
     BaseModel,
     BeforeValidator,
-    ConfigDict,
     Field,
     field_validator,
     model_serializer,
@@ -730,11 +729,9 @@ def _sanitise_name(v: str) -> str:
     return re.sub(r"[^a-zA-Z0-9_-]", "_", v).strip()
 
 
-class ChatEntry(BaseModel):
+class ChatEntry(BaseModel, use_enum_values=True):
     """Represents a message in the chat context."""
 
-    model_config = ConfigDict(use_enum_values=True)
-
     role: ChatRole = Field(
         description="Who said the message?",
     )

docs/pgaudit_guide.md

@@ -39,7 +39,7 @@ spec:
 - customize audit_statement based on the level of DML statement you would want to monitor
 
 ```python
-async def _grant_auditor_priviledge(engine: AsyncEngine) -> bool:
+async def _grant_auditor_privilege(engine: AsyncEngine) -> bool:
     """
     Apply the necessary grants to allow the auditor role to audit the database.
     """

services/api/pyproject.toml

@@ -75,7 +75,7 @@ unfixable = ["B"]
 "**/{tests,docs,tools}/*" = ["E402"]
 
 [tool.ruff.lint.isort]
-known-first-party = ["jamaibase", "owl"]
+known-first-party = ["jamaibase", "owl", "polylm"]
 
 [tool.ruff.lint.flake8-bugbear]
 # Allow default arguments like, e.g., `data: List[str] = fastapi.Query(None)`.
@@ -131,7 +131,7 @@ dependencies = [
     "fastapi[standard]~=0.115.0",
     "flower~=2.0.0",
     "gunicorn~=22.0.0",
-    "httpx~=0.27.0",
+    "httpx~=0.27",
     "itsdangerous~=2.2.0",
     "jamaibase>=0.4.1",
     "lancedb==0.12.0",
@@ -208,8 +208,7 @@ dependencies = [
     "typing_extensions~=4.14.0",
     "uuid-utils~=0.9.0",
     "uuid7~=0.1.0",
-    # uvicorn 0.29.x shutdown seems unclean and 0.30.x child process sometimes dies
-    "uvicorn[standard]~=0.28.0",
+    "uvicorn[standard]~=0.28",
     "xmltodict~=0.14.0",
 ]
 dynamic = ["version"]

services/api/src/owl/scripts/backup_db.py services/api/scripts/backup_db.pyservices/api/src/owl/scripts/backup_db.py renamed to services/api/scripts/backup_db.py

@@ -27,7 +27,6 @@ class EnvConfig(BaseSettings):
     max_concurrency: int = 300
     db_init: bool | None = None  # None means unset
     db_reset: bool = False
-    db_init_max_users: int = 5
     cache_reset: bool = False
     enable_byok: bool = True
     disable_billing: bool = False
@@ -139,31 +138,25 @@ def check_alternate_service_key(self) -> Self:
 
     @model_validator(mode="after")
     def validate_db_path(self) -> Self:
-        """
-        Validates that `db_path` starts with either `rqlite+pyrqlite://` or `sqlite://` or `sqlite+libsql://` or `postgresql`.
-        """
         if not (
-            self.db_path.startswith("rqlite+pyrqlite://")
-            or self.db_path.startswith("sqlite://")
-            or self.db_path.startswith("sqlite+libsql://")
-            or self.db_path.startswith("postgresql")
+            self.db_path.startswith("postgresql+psycopg://")
+            # or self.db_path.startswith("sqlite+aiosqlite://")
+            # or self.db_path.startswith("sqlite://")
+            # or self.db_path.startswith("sqlite+libsql://")
+            # or self.db_path.startswith("rqlite+pyrqlite://")
         ):
             raise ValueError(f'`db_path` "{self.db_path}" has an invalid dialect.')
         return self
 
     @property
-    def db_dialect(self) -> Literal["rqlite", "libsql", "postgresql", "sqlite"]:
+    def db_dialect(self) -> Literal["postgresql"]:
         """
-        Show the sqlite dialect that's in use based on the `db_path`.
+        Show the dialect that's in use based on the `db_path`.
         """
-        if self.db_path.startswith("rqlite+pyrqlite://"):
-            return "rqlite"
-        elif self.db_path.startswith("sqlite+libsql://"):
-            return "libsql"
-        elif self.db_path.startswith("postgresql"):
+        if self.db_path.startswith("postgresql"):
             return "postgresql"
-        elif self.db_path.startswith("sqlite://"):
-            return "sqlite"
+        else:
+            raise ValueError(f'`db_path` "{self.db_path}" has an invalid dialect.')
 
     @cached_property
     def is_oss(self) -> bool:

…ces/api/src/owl/scripts/create_meters.py services/api/scripts/create_meters.pyservices/api/src/owl/scripts/create_meters.py renamed to services/api/scripts/create_meters.py services/api/src/owl/scripts/create_meters.py renamed to services/api/scripts/create_meters.py

@@ -97,37 +97,12 @@ def cached_text(query: str) -> TextClause:
     return text(query)
 
 
-async def reset_db(*, reset_max_users: int = 3):
-    from sqlmodel import func, select
-
-    from owl.db.models import User
-
+async def reset_db():
     # Only allow DB reset in dev with localhost
     if "@localhost:" not in ENV_CONFIG.db_path:
         raise ValueError("DB reset is only allowed in dev with localhost DB.")
 
     async with async_session() as session:
-        # As a safety measure, reset DB only if it has less than `init_max_users` users
-        # Just in case we accidentally tried to nuke a prod DB
-        user_table_exists = (
-            await session.exec(
-                text(
-                    (
-                        f"SELECT EXISTS ("
-                        f"SELECT FROM information_schema.tables WHERE table_schema = '{SCHEMA}' AND table_name = 'User'"
-                        ");"
-                    )
-                )
-            )
-        ).scalar()
-        if user_table_exists:
-            user_count = (await session.exec(select(func.count(User.id)))).one()
-            if user_count >= reset_max_users:
-                logger.info(
-                    f"Found {user_count:,d} users, abort database reset (>= {reset_max_users} users)."
-                )
-                return
-
         # Delete all tables
         logger.warning(f'Resetting database (dropping schema "{SCHEMA}")...')
         await session.exec(text(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE"))
@@ -702,7 +677,7 @@ async def migrate_db():
     await CACHE.aclose()
 
 
-async def init_db(*, init_max_users: int = 3):
+async def init_db():
     from fastapi import Request
     from sqlmodel import func, select
     from starlette.datastructures import URL, Headers
@@ -721,19 +696,16 @@ async def init_db(*, init_max_users: int = 3):
     )
 
     async with async_session() as session:
-        # As a safety measure, init DB only if it has less than `init_max_users` users
-        # Just in case we accidentally tried to nuke a prod DB
         user_count = (await session.exec(select(func.count(User.id)))).one()
-        if user_count >= init_max_users:
-            logger.info(
-                f"Found {user_count:,d} users, abort database initialisation (>= {init_max_users} users)."
-            )
-            return
-
         # Only enforce OSS check if db_init=False
-        if ENV_CONFIG.is_oss and user_count != 0:
-            logger.info("OSS mode: Skipping initialization (non-empty DB).")
-            return
+        if ENV_CONFIG.is_oss:
+            if user_count != 0:
+                logger.info("OSS mode: Skipping initialization (non-empty DB).")
+                return
+        else:
+            # Only allow DB init in dev with localhost
+            if "@localhost:" not in ENV_CONFIG.db_path:
+                raise ValueError("DB init is only allowed in dev with localhost DB.")
 
         logger.info("Initialising database...")
 

services/api/src/owl/configs/oss.py

@@ -2423,6 +2423,7 @@ def _measure_ram() -> str:
                 logger.info(f'Importing table "{self.table_id}": Skipped S3 upload.')
         prog.parse_data.progress = 100
         await CACHE.set_progress(prog)
+        semaphore = Semaphore(S3_MAX_CONCURRENCY)
 
         async def _upload(
             old_uri: str,
@@ -2441,7 +2442,6 @@ async def _upload(
                 return (old_uri, new_uri)
 
         uris_seen: dict[str, str] = {}  # Old URI to new URI
-        semaphore = Semaphore(S3_MAX_CONCURRENCY)
         upload_coros = []
         for row in rows:
             file_byte_cols = [c for c in row.keys() if c.endswith("__")]