From f7df29792f93a9d9e76bb474ff15dce7d688b30e Mon Sep 17 00:00:00 2001
From: Ruthenic <mdrakea3@tutanota.com>
Date: Sun, 16 Oct 2022 01:55:05 -0400
Subject: [PATCH] properly do token rotation (for real this time)

---
 src/ai/huggingface.ts | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/ai/huggingface.ts b/src/ai/huggingface.ts
index cbc9219..4a3dbdf 100644
--- a/src/ai/huggingface.ts
+++ b/src/ai/huggingface.ts
@@ -36,7 +36,8 @@ class HuggingFaceAI {
         this.memory = db.history ? db.history.map((m) => `${m.name}: "${m.content}"`) : [];
     }
 
-    async #query(prompt: string) {
+    //TODO: type our response object
+    async #query(prompt: string): Promise<Record<string, string>[]> {
         const res = await fetch(`https://api-inference.huggingface.co/models/${this.model}`, {
             body: JSON.stringify({
                 inputs: prompt,
@@ -51,7 +52,14 @@ class HuggingFaceAI {
         if (this.tokenNum > this.tokens.length - 1) {
             this.tokenNum = 0
         }
-        return await res.json();
+
+        const jsonRes = await res.json();
+        if (!jsonRes[0].generated_text) {
+            console.log(jsonRes)
+            console.warn("Retrying generation with new token...")
+            return await this.#query(prompt) // bound to go well 
+        }
+        return jsonRes
     }
 
     reset() {