🥅 Catch SSE error (huggingface#169)

Author: coyotte508

packages/inference/src/tasks/custom/streamingRequest.ts

@@ -65,7 +65,11 @@ export async function* streamingRequest<T>(
 			onChunk(value);
 			for (const event of events) {
 				if (event.data.length > 0) {
-					yield JSON.parse(event.data) as T;
+					const data = JSON.parse(event.data);
+					if (typeof data === "object" && data !== null && "error" in data) {
+						throw new Error(data.error);
+					}
+					yield data as T;
 				}
 			}
 			events = [];

packages/inference/src/tasks/nlp/textGeneration.ts

@@ -44,6 +44,10 @@ export type TextGenerationArgs = BaseArgs & {
 		 * (Default: None). Float to define the tokens that are within the sample operation of text generation. Add tokens in the sample for more probable to least probable until the sum of the probabilities is greater than top_p.
 		 */
 		top_p?: number;
+		/**
+		 * (Default: None). Integer. The maximum number of tokens from the input.
+		 */
+		truncate?: number;
 	};
 };
 

packages/inference/test/HfInference.spec.ts

@@ -216,6 +216,20 @@ describe.concurrent(
 			}
 		});
 
+		it("textGenerationStream - catch error", async () => {
+			const response = hf.textGenerationStream({
+				model: "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+				inputs: "Write a short story about a robot that becomes sentient and takes over the world.",
+				parameters: {
+					truncate: 1024,
+				},
+			});
+
+			await expect(response.next()).rejects.toThrow(
+				"Input validation error: `truncate` must be strictly positive and less than 1000. Given: 1024"
+			);
+		});
+
 		it("tokenClassification", async () => {
 			expect(
 				await hf.tokenClassification({

packages/inference/test/tapes.json

@@ -628,5 +628,26 @@
         "vary": "Accept-Encoding, Origin, Access-Control-Request-Method, Access-Control-Request-Headers"
       }
     }
+  },
+  "1841a1a3806a917b0b01a7e066c425d4bd8f2da39168ed144061299c9a277b34": {
+    "url": "https://api-inference.huggingface.co/models/OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+    "init": {
+      "headers": {
+        "Content-Type": "application/json"
+      },
+      "method": "POST",
+      "body": "{\"inputs\":\"Write a short story about a robot that becomes sentient and takes over the world.\",\"parameters\":{\"truncate\":1024},\"stream\":true,\"options\":{}}"
+    },
+    "response": {
+      "body": "data:{\"error\":\"Input validation error: `truncate` must be strictly positive and less than 1000. Given: 1024\",\"error_type\":\"validation\"}\n\n",
+      "status": 200,
+      "statusText": "OK",
+      "headers": {
+        "access-control-allow-credentials": "true",
+        "connection": "keep-alive",
+        "content-type": "text/event-stream",
+        "vary": "Origin, Access-Control-Request-Method, Access-Control-Request-Headers"
+      }
+    }
   }
 }