ENG-1129: aixplain sdk caching functions

aixplain/enums/__init__.py

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

aixplain/enums/cache_utils.py

@@ -0,0 +1,29 @@
+import os
+import json
+import time
+import logging
+
+CACHE_DURATION = 24 * 60 * 60 
+
+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):
+    try:
+        with open(cache_file, "r") as f:
+            cache_data = json.load(f)
+            if time.time() - cache_data["timestamp"] < CACHE_DURATION:
+                logging.info(f"Loaded valid cache from {cache_file}.")
+                return cache_data["data"]
+            else:
+                logging.info(f"Cache expired for {cache_file}.")
+                return None
+    except FileNotFoundError:
+        return None
+    except json.JSONDecodeError:
+        return None

aixplain/enums/function.py

@@ -27,34 +27,51 @@
 from aixplain.utils.request_utils import _request_with_retry
 from enum import Enum
 from urllib.parse import urljoin
+import logging
+from .cache_utils import save_to_cache, load_from_cache
 
+CACHE_FILE = ".aixplain_cache/functions.json"
 
 def load_functions():
-    api_key = config.TEAM_API_KEY
-    backend_url = config.BACKEND_URL
-
-    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()
-    functions = Enum("Function", {w["id"].upper().replace("-", "_"): w["id"] for w in resp["items"]}, type=str)
-    functions_input_output = {
-        function["id"]: {
-            "input": {
-                input_data_object["dataType"]
-                for input_data_object in function["params"]
-                if input_data_object["required"] is True
-            },
-            "output": {output_data_object["dataType"] for output_data_object in function["output"]},
-            "spec": function,
+    cached_data = load_from_cache(CACHE_FILE)
+    if cached_data:
+        return Enum("Function", cached_data["enum"], type=str), cached_data["input_output"]
+
+
+    try:
+        api_key = config.TEAM_API_KEY
+        backend_url = config.BACKEND_URL
+        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("Functions could not be loaded. Invalid API key or server issue.")
+
+        resp = r.json()
+        functions_enum = {
+            w["id"].upper().replace("-", "_"): w["id"] for w in resp["items"]
         }
-        for function in resp["items"]
-    }
-    return functions, functions_input_output
+        functions_input_output = {
+            function["id"]: {
+                "input": {
+                    input_data_object["dataType"]
+                    for input_data_object in function["params"]
+                    if input_data_object["required"] is True
+                },
+                "output": {output_data_object["dataType"] for output_data_object in function["output"]},
+                "spec": function,
+            }
+            for function in resp["items"]
+        }
+
+        save_to_cache(CACHE_FILE, {"enum": functions_enum, "input_output": functions_input_output})
+
+        return Enum("Function", functions_enum, type=str), functions_input_output
+
+    except Exception as e:
+        logging.error(f"Failed to load functions from API: {e}")
+        raise Exception("Unable to load functions from cache or API.")
+
 
 Function, FunctionInputOutput = load_functions()

aixplain/enums/language.py

@@ -21,15 +21,19 @@
     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 .cache_utils import save_to_cache, load_from_cache
 
+CACHE_FILE = ".aixplain_cache/languages.json"
 
 def load_languages():
+    cached_languages = load_from_cache(CACHE_FILE)
+    if cached_languages:
+        return Enum("Language", cached_languages, type=dict)
+
     api_key = config.TEAM_API_KEY
     backend_url = config.BACKEND_URL
 
@@ -52,6 +56,7 @@ def load_languages():
             dialect_value = dialect["value"]
 
             languages[language_label + "_" + dialect_label] = {"language": language, "dialect": dialect_value}
+    save_to_cache(CACHE_FILE, languages)
     return Enum("Language", languages, type=dict)
 
 

aixplain/enums/license.py

@@ -22,14 +22,19 @@
 """
 
 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 .cache_utils import save_to_cache, load_from_cache
 
+CACHE_FILE = ".aixplain_cache/licenses.json"
 
 def load_licenses():
+    cached_licenses = load_from_cache(CACHE_FILE)
+    if cached_licenses:
+        return Enum("License", cached_licenses, type=str)
+
     try:
         api_key = config.TEAM_API_KEY
         backend_url = config.BACKEND_URL
@@ -43,7 +48,10 @@ def load_licenses():
                 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)
+        licenses = {"_".join(w["name"].split()): w["id"] for w in resp}
+
+        save_to_cache(CACHE_FILE, licenses)
+        return Enum("License", licenses, type=str)
     except Exception as e:
         logging.exception("License Loading Error")
         raise Exception("License Loading Error")