Merge pull request #49 from FocoosAI/feat/add-user-info

feat: add user and quota models to enhance user info retrieval 

Introduce `User` and `Quotas` models to provide structured user information, including API key and usage quotas. This change improves the clarity and usability of user data, facilitating better resource management and integration.

.pre-commit-config.yaml

@@ -23,30 +23,30 @@ repos:
   #     - id: mypy
   #       args: [--ignore-missing-imports, --pretty, --show-error-codes]
   - repo: https://github.com/nbQA-dev/nbQA # run lint on notebook
-    rev: 1.5.3
+    rev: 1.9.1
     hooks:
       - id: nbqa-black
       - id: nbqa-pyupgrade
     #  - id: nbqa-isort
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.7.0
+    rev: v3.19.1
     hooks:
       - id: pyupgrade
   - repo: https://github.com/pycqa/isort
-    rev: 5.12.0
+    rev: 5.13.2
     hooks:
       - id: isort
         name: isort (python)
         args: [--profile, black]
   - repo: https://github.com/psf/black
-    rev: 23.3.0
+    rev: 24.10.0
     hooks:
       - id: black
   - repo: https://github.com/kynan/nbstripout
     rev: 0.6.0
     hooks:
       - id: nbstripout
   - repo: https://github.com/hadialqattan/pycln # remove unused import
-    rev: v2.4.0
+    rev: v2.5.0
     hooks:
       - id: pycln

focoos/focoos.py

@@ -20,7 +20,14 @@
 
 from focoos.config import FOCOOS_CONFIG
 from focoos.local_model import LocalModel
-from focoos.ports import DatasetMetadata, ModelMetadata, ModelPreview, RuntimeTypes
+from focoos.ports import (
+    DatasetMetadata,
+    ModelMetadata,
+    ModelPreview,
+    Quotas,
+    RuntimeTypes,
+    User,
+)
 from focoos.remote_model import RemoteModel
 from focoos.utils.logger import setup_logging
 from focoos.utils.system import HttpClient
@@ -83,13 +90,13 @@ def __init__(
         host_url = host_url or FOCOOS_CONFIG.default_host_url
 
         self.http_client = HttpClient(self.api_key, host_url)
-        self.user_info = self._get_user_info()
+        self.user_info = self.get_user_info()
         self.cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "focoos")
         logger.info(
-            f"Currently logged as: {self.user_info['email']} environment: {host_url}"
+            f"Currently logged as: {self.user_info.email} environment: {host_url}"
         )
 
-    def _get_user_info(self):
+    def get_user_info(self) -> User:
         """
         Retrieves information about the authenticated user.
 
@@ -103,7 +110,7 @@ def _get_user_info(self):
         if res.status_code != 200:
             logger.error(f"Failed to get user info: {res.status_code} {res.text}")
             raise ValueError(f"Failed to get user info: {res.status_code} {res.text}")
-        return res.json()
+        return User.from_json(res.json())
 
     def get_model_info(self, model_name: str) -> ModelMetadata:
         """

focoos/ports.py

@@ -307,3 +307,32 @@ def pretty_print(self):
             else:
                 print(f"{key}: {value}")
         print("================================================")
+
+
+class ApiKey(FocoosBaseModel):
+    key: str  # type: ignore
+
+
+class Quotas(FocoosBaseModel):
+    # INFERENCE
+    total_inferences: int
+    max_inferences: int
+    # STORAGE
+    used_storage_gb: float
+    max_storage_gb: float
+    # TRAINING
+    active_training_jobs: list[str]
+    max_active_training_jobs: int
+
+    # ML_G4DN_XLARGE TRAINING HOURS
+    used_mlg4dnxlarge_training_jobs_hours: float
+    max_mlg4dnxlarge_training_jobs_hours: float
+
+
+class User(FocoosBaseModel):
+    email: str
+    created_at: datetime
+    updated_at: datetime
+    company: Optional[str] = None
+    api_key: ApiKey
+    quotas: Quotas

notebooks/playground.ipynb

@@ -439,6 +439,29 @@
     "pprint(logs)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Get User Info"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from focoos import Focoos, DEV_API_URL\n",
+    "import os\n",
+    "from pprint import pprint\n",
+    "\n",
+    "focoos = Focoos(api_key=os.getenv(\"FOCOOS_API_KEY\"), host_url=DEV_API_URL)\n",
+    "\n",
+    "user_info = focoos.get_user_info()\n",
+    "pprint(user_info.model_dump())"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -457,6 +480,13 @@
     "system_info = get_system_info()\n",
     "system_info.pretty_print()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -475,7 +505,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.15"
+   "version": "3.12.0"
   }
  },
  "nbformat": 4,

pyproject.toml

@@ -16,7 +16,7 @@ dependencies = [
     "Pillow~=10.4.0",
     "supervision~=0.26.0rc2",
     "pydantic~=2.8.2",
-    "pydantic-settings~=2.5.0",
+    "pydantic-settings~=2.7.1",
     "tqdm~=4.67.1",
     "numpy~=2.2.1",
     "scipy~=1.14.1",

tests/test_focoos.py

@@ -1,5 +1,6 @@
 import pathlib
 import tempfile
+from datetime import datetime
 from unittest.mock import MagicMock
 
 import pytest
@@ -16,7 +17,23 @@
 def focoos_instance(mock_http_client) -> Focoos:
     """Fixture to provide a Focoos instance with a mocked HttpClient."""
     mock_http_client.get.return_value.status_code = 200
-    mock_http_client.get.return_value.json.return_value = {"email": "[email protected]"}
+    mock_http_client.get.return_value.json.return_value = {
+        "email": "[email protected]",
+        "created_at": "2024-01-01",
+        "updated_at": "2025-01-01",
+        "company": "test_company",
+        "api_key": {"key": "test_api_key"},
+        "quotas": {
+            "total_inferences": 10,
+            "max_inferences": 1000,
+            "used_storage_gb": 10,
+            "max_storage_gb": 1000,
+            "active_training_jobs": ["job1"],
+            "max_active_training_jobs": 1,
+            "used_mlg4dnxlarge_training_jobs_hours": 10,
+            "max_mlg4dnxlarge_training_jobs_hours": 1000,
+        },
+    }
     return Focoos(api_key="test_api_key", host_url="http://mock-host-url.com")
 
 
@@ -98,7 +115,18 @@ def test_focoos_initialization_fail_to_fetch_user_info(focoos_instance: Focoos):
 
 def test_focoos_initialization(focoos_instance: Focoos):
     assert focoos_instance.api_key == "test_api_key"
-    assert focoos_instance.user_info["email"] == "[email protected]"
+    assert focoos_instance.user_info.email == "[email protected]"
+    assert focoos_instance.user_info.company == "test_company"
+    assert focoos_instance.user_info.created_at == datetime(2024, 1, 1)
+    assert focoos_instance.user_info.updated_at == datetime(2025, 1, 1)
+    assert focoos_instance.user_info.quotas.total_inferences == 10
+    assert focoos_instance.user_info.quotas.max_inferences == 1000
+    assert focoos_instance.user_info.quotas.used_storage_gb == 10
+    assert focoos_instance.user_info.quotas.max_storage_gb == 1000
+    assert focoos_instance.user_info.quotas.active_training_jobs == ["job1"]
+    assert focoos_instance.user_info.quotas.max_active_training_jobs == 1
+    assert focoos_instance.user_info.quotas.used_mlg4dnxlarge_training_jobs_hours == 10
+    assert focoos_instance.user_info.quotas.max_mlg4dnxlarge_training_jobs_hours == 1000
 
 
 def test_get_model_info(focoos_instance: Focoos):