ENG-1129: aixplain sdk caching functions

aixplain/enums/__init__.py

@@ -13,4 +13,4 @@
 from .supplier import Supplier
 from .sort_by import SortBy
 from .sort_order import SortOrder
-from .response_status import ResponseStatus
+from .response_status import ResponseStatus

aixplain/enums/function.py

@@ -21,27 +21,32 @@
     Function Enum
 """
 
-import logging
-
 from aixplain.utils import config
 from aixplain.utils.request_utils import _request_with_retry
 from enum import Enum
 from urllib.parse import urljoin
+from aixplain.utils.cache_utils import save_to_cache, load_from_cache, CACHE_FOLDER
+
+CACHE_FILE = f"{CACHE_FOLDER}/functions.json"
 
 
 def load_functions():
     api_key = config.TEAM_API_KEY
     backend_url = config.BACKEND_URL
 
-    url = urljoin(backend_url, "sdk/functions")
+    resp = load_from_cache(CACHE_FILE)
+    if resp is None:
+        url = urljoin(backend_url, "sdk/functions")
+
+        headers = {"x-api-key": api_key, "Content-Type": "application/json"}
+        r = _request_with_retry("get", url, headers=headers)
+        if not 200 <= r.status_code < 300:
+            raise Exception(
+                f'Functions could not be loaded, probably due to the set API key (e.g. "{api_key}") is not valid. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)'
+            )
+        resp = r.json()
+        save_to_cache(CACHE_FILE, resp)
 
-    headers = {"x-api-key": api_key, "Content-Type": "application/json"}
-    r = _request_with_retry("get", url, headers=headers)
-    if not 200 <= r.status_code < 300:
-        raise Exception(
-            f'Functions could not be loaded, probably due to the set API key (e.g. "{api_key}") is not valid. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)'
-        )
-    resp = r.json()
     functions = Enum("Function", {w["id"].upper().replace("-", "_"): w["id"] for w in resp["items"]}, type=str)
     functions_input_output = {
         function["id"]: {
@@ -57,4 +62,5 @@ def load_functions():
     }
     return functions, functions_input_output
 
+
 Function, FunctionInputOutput = load_functions()

aixplain/enums/language.py

@@ -21,27 +21,32 @@
     Language Enum
 """
 
-import logging
-
-from aixplain.utils import config
-from aixplain.utils.request_utils import _request_with_retry
 from enum import Enum
 from urllib.parse import urljoin
+from aixplain.utils import config
+from aixplain.utils.request_utils import _request_with_retry
+from aixplain.utils.cache_utils import save_to_cache, load_from_cache, CACHE_FOLDER
+
+CACHE_FILE = f"{CACHE_FOLDER}/languages.json"
 
 
 def load_languages():
-    api_key = config.TEAM_API_KEY
-    backend_url = config.BACKEND_URL
-
-    url = urljoin(backend_url, "sdk/languages")
-
-    headers = {"x-api-key": api_key, "Content-Type": "application/json"}
-    r = _request_with_retry("get", url, headers=headers)
-    if not 200 <= r.status_code < 300:
-        raise Exception(
-            f'Languages could not be loaded, probably due to the set API key (e.g. "{api_key}") is not valid. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)'
-        )
-    resp = r.json()
+    resp = load_from_cache(CACHE_FILE)
+    if resp is None:
+        api_key = config.TEAM_API_KEY
+        backend_url = config.BACKEND_URL
+
+        url = urljoin(backend_url, "sdk/languages")
+
+        headers = {"x-api-key": api_key, "Content-Type": "application/json"}
+        r = _request_with_retry("get", url, headers=headers)
+        if not 200 <= r.status_code < 300:
+            raise Exception(
+                f'Languages could not be loaded, probably due to the set API key (e.g. "{api_key}") is not valid. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)'
+            )
+        resp = r.json()
+        save_to_cache(CACHE_FILE, resp)
+
     languages = {}
     for w in resp:
         language = w["value"]

aixplain/enums/license.py

@@ -22,29 +22,36 @@
 """
 
 import logging
-
-from aixplain.utils import config
-from aixplain.utils.request_utils import _request_with_retry
 from enum import Enum
 from urllib.parse import urljoin
+from aixplain.utils import config
+from aixplain.utils.request_utils import _request_with_retry
+from aixplain.utils.cache_utils import save_to_cache, load_from_cache, CACHE_FOLDER
+
+CACHE_FILE = f"{CACHE_FOLDER}/licenses.json"
 
 
 def load_licenses():
+    resp = load_from_cache(CACHE_FILE)
+
     try:
-        api_key = config.TEAM_API_KEY
-        backend_url = config.BACKEND_URL
-
-        url = urljoin(backend_url, "sdk/licenses")
-
-        headers = {"x-api-key": api_key, "Content-Type": "application/json"}
-        r = _request_with_retry("get", url, headers=headers)
-        if not 200 <= r.status_code < 300:
-            raise Exception(
-                f'Licenses could not be loaded, probably due to the set API key (e.g. "{api_key}") is not valid. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)'
-            )
-        resp = r.json()
-        return Enum("License", {"_".join(w["name"].split()): w["id"] for w in resp}, type=str)
-    except Exception as e:
+        if resp is None:
+            api_key = config.TEAM_API_KEY
+            backend_url = config.BACKEND_URL
+
+            url = urljoin(backend_url, "sdk/licenses")
+
+            headers = {"x-api-key": api_key, "Content-Type": "application/json"}
+            r = _request_with_retry("get", url, headers=headers)
+            if not 200 <= r.status_code < 300:
+                raise Exception(
+                    f'Licenses could not be loaded, probably due to the set API key (e.g. "{api_key}") is not valid. For help, please refer to the documentation (https://github.com/aixplain/aixplain#api-key-setup)'
+                )
+            resp = r.json()
+            save_to_cache(CACHE_FILE, resp)
+        licenses = {"_".join(w["name"].split()): w["id"] for w in resp}
+        return Enum("License", licenses, type=str)
+    except Exception:
         logging.exception("License Loading Error")
         raise Exception("License Loading Error")
 

aixplain/utils/cache_utils.py

@@ -0,0 +1,27 @@
+import os
+import json
+import time
+import logging
+
+CACHE_DURATION = 24 * 60 * 60
+CACHE_FOLDER = ".aixplain_cache"
+
+
+def save_to_cache(cache_file, data):
+    try:
+        os.makedirs(os.path.dirname(cache_file), exist_ok=True)
+        with open(cache_file, "w") as f:
+            json.dump({"timestamp": time.time(), "data": data}, f)
+    except Exception as e:
+        logging.error(f"Failed to save cache to {cache_file}: {e}")
+
+
+def load_from_cache(cache_file):
+    if os.path.exists(cache_file) is True:
+        with open(cache_file, "r") as f:
+            cache_data = json.load(f)
+            if time.time() - cache_data["timestamp"] < CACHE_DURATION:
+                return cache_data["data"]
+            else:
+                return None
+    return None

aixplain/utils/config.py

@@ -26,7 +26,9 @@
 
 if AIXPLAIN_API_KEY and TEAM_API_KEY:
     if AIXPLAIN_API_KEY != TEAM_API_KEY:
-        raise Exception("Conflicting API keys: 'AIXPLAIN_API_KEY' and 'TEAM_API_KEY' are both provided but do not match. Please provide only one API key.")
+        raise Exception(
+            "Conflicting API keys: 'AIXPLAIN_API_KEY' and 'TEAM_API_KEY' are both provided but do not match. Please provide only one API key."
+        )
 
 
 if AIXPLAIN_API_KEY and not TEAM_API_KEY: