From 32c4589df3b0fa86619922e541c0309adb2c53f5 Mon Sep 17 00:00:00 2001
From: DominikDoom <rehdominik@gmail.com>
Date: Tue, 29 Aug 2023 09:39:32 +0200
Subject: [PATCH] Rework wildcards to use own API endpoint Maybe fixes #226

---
 javascript/ext_wildcards.js        | 16 ++++++++++------
 scripts/tag_autocomplete_helper.py | 24 ++++++++++++++++++++----
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/javascript/ext_wildcards.js b/javascript/ext_wildcards.js
index 4eb97c6..328afdb 100644
--- a/javascript/ext_wildcards.js
+++ b/javascript/ext_wildcards.js
@@ -22,10 +22,13 @@ class WildcardParser extends BaseTagParser {
     
         let wildcards = [];
         for (let i = 0; i < wcPairs.length; i++) {
-            const wcPair = wcPairs[i];
-            if (!wcPair[0] || !wcPair[1]) continue;
+            const basePath = wcPairs[i][0];
+            const fileName = wcPairs[i][1];
+            if (!basePath || !fileName) return;
 
-            if (wcPair[0].endsWith(".yaml")) {
+            // YAML wildcards are already loaded as json, so we can get the values directly.
+            // basePath is the name of the file in this case, and fileName the key
+            if (basePath.endsWith(".yaml")) {
                 const getDescendantProp = (obj, desc) => {
                     const arr = desc.split("/");
                     while (arr.length) {
@@ -33,10 +36,11 @@ class WildcardParser extends BaseTagParser {
                     }
                     return obj;
                 }
-                wildcards = wildcards.concat(getDescendantProp(yamlWildcards[wcPair[0]], wcPair[1]));
+                wildcards = wildcards.concat(getDescendantProp(yamlWildcards[basePath], fileName));
             } else {
-                const fileContent = (await readFile(`${wcPair[0]}${wcPair[1]}.txt`)).split("\n")
-                .filter(x => x.trim().length > 0 && !x.startsWith('#'));  // Remove empty lines and comments
+                const fileContent = (await fetchAPI(`tacapi/v1/wildcard-contents?basepath=${basePath}&filename=${fileName}.txt`, false))
+                    .split("\n")
+                    .filter(x => x.trim().length > 0 && !x.startsWith('#'));  // Remove empty lines and comments
                 wildcards = wildcards.concat(fileContent);
             }
         }
diff --git a/scripts/tag_autocomplete_helper.py b/scripts/tag_autocomplete_helper.py
index 3002058..b735701 100644
--- a/scripts/tag_autocomplete_helper.py
+++ b/scripts/tag_autocomplete_helper.py
@@ -457,18 +457,18 @@ script_callbacks.on_ui_settings(on_ui_settings)
 def api_tac(_: gr.Blocks, app: FastAPI):
     async def get_json_info(base_path: Path, filename: str = None):
         if base_path is None or (not base_path.exists()):
-            return json.dumps({})
+            return JSONResponse({}, status_code=404)
         
         try:
             json_candidates = glob.glob(base_path.as_posix() + f"/**/{filename}.json", recursive=True)
             if json_candidates is not None and len(json_candidates) > 0:
                 return FileResponse(json_candidates[0])
         except Exception as e:
-            return json.dumps({"error": e})
+            return JSONResponse({"error": e}, status_code=500)
         
     async def get_preview_thumbnail(base_path: Path, filename: str = None, blob: bool = False):
         if base_path is None or (not base_path.exists()):
-            return json.dumps({})
+            return JSONResponse({}, status_code=404)
         
         try:
             img_glob = glob.glob(base_path.as_posix() + f"/**/{filename}.*", recursive=True)
@@ -479,7 +479,7 @@ def api_tac(_: gr.Blocks, app: FastAPI):
                 else:
                     return JSONResponse({"url": urllib.parse.quote(img_candidates[0])})
         except Exception as e:
-            return json.dumps({"error": e})
+            return JSONResponse({"error": e}, status_code=500)
 
     @app.get("/tacapi/v1/lora-info/{lora_name}")
     async def get_lora_info(lora_name):
@@ -509,7 +509,23 @@ def api_tac(_: gr.Blocks, app: FastAPI):
     async def get_thumb_preview_blob(filename, type):
         return await get_preview_thumbnail(get_path_for_type(type), filename, True)
         
+    @app.get("/tacapi/v1/wildcard-contents")
+    async def get_wildcard_contents(basepath: str, filename: str):
+        if basepath is None or basepath == "":
+            return JSONResponse({}, status_code=404)
 
+        base = Path(basepath)
+        if base is None or (not base.exists()):
+            return JSONResponse({}, status_code=404)
+        
+        try:
+            wildcard_path = base.joinpath(filename)
+            if wildcard_path.exists():
+                return FileResponse(wildcard_path)
+            else:
+                return JSONResponse({}, status_code=404)
+        except Exception as e:
+            return JSONResponse({"error": e}, status_code=500)
 
 
 script_callbacks.on_app_started(api_tac)