💚 Fix VCR module (#133)

Author: coyotte508

packages/inference/src/HfInference.ts

@@ -1048,7 +1048,9 @@ export class HfInference {
 			throw new Error(`Server response contains error: ${response.status}`);
 		}
 		if (response.headers.get("content-type") !== "text/event-stream") {
-			throw new Error(`Server does not support event stream content type`);
+			throw new Error(
+				`Server does not support event stream content type, it returned ` + response.headers.get("content-type")
+			);
 		}
 
 		const reader = response.body.getReader();

packages/inference/src/utils/base64FromBytes.ts

@@ -0,0 +1,11 @@
+export function base64FromBytes(arr: Uint8Array): string {
+	if (globalThis.Buffer) {
+		return globalThis.Buffer.from(arr).toString("base64");
+	} else {
+		const bin: string[] = [];
+		arr.forEach((byte) => {
+			bin.push(String.fromCharCode(byte));
+		});
+		return globalThis.btoa(bin.join(""));
+	}
+}

packages/inference/src/utils/omit.ts

@@ -0,0 +1,11 @@
+import { pick } from "./pick";
+import { typedInclude } from "./typed-include";
+
+/**
+ * Return copy of object, omitting blocklisted array of props
+ */
+export function omit<T, K extends keyof T>(o: T, props: K[] | K): Pick<T, Exclude<keyof T, K>> {
+	const propsArr = Array.isArray(props) ? props : [props];
+	const letsKeep = (Object.keys(o) as (keyof T)[]).filter((prop) => !typedInclude(propsArr, prop));
+	return pick(o, letsKeep);
+}

packages/inference/src/utils/pick.ts

@@ -0,0 +1,16 @@
+/**
+ * Return copy of object, only keeping allowlisted properties.
+ *
+ * This doesn't add {p: undefined} anymore, for props not in the o object.
+ */
+export function pick<T, K extends keyof T>(o: T, props: K[] | ReadonlyArray<K>): Pick<T, K> {
+	// inspired by stackoverflow.com/questions/25553910/one-liner-to-take-some-properties-from-object-in-es-6
+	return Object.assign(
+		{},
+		...props.map((prop) => {
+			if (o[prop] !== undefined) {
+				return { [prop]: o[prop] };
+			}
+		})
+	);
+}

packages/inference/src/utils/typed-include.ts

@@ -0,0 +1,3 @@
+export function typedInclude<V, T extends V>(arr: readonly T[], v: V): v is T {
+	return arr.includes(v as T);
+}

packages/inference/test/tapes.json

@@ -1,6 +1,9 @@
+import { base64FromBytes } from "../src/utils/base64FromBytes";
 import { isBackend, isFrontend } from "../src/utils/env-predicates";
+import { omit } from "../src/utils/omit";
 
 const TAPES_FILE = "./tapes.json";
+const BASE64_PREFIX = "data:application/octet-stream;base64,";
 
 enum MODE {
 	RECORD = "record",
@@ -28,9 +31,7 @@ if (process.env.VCR_MODE) {
 
 const originalFetch = globalThis.fetch;
 
-globalThis.fetch = (...args) => {
-	return vcr(originalFetch, args[0], args[1]);
-};
+globalThis.fetch = (...args) => vcr(originalFetch, args[0], args[1]);
 
 /**
  * Represents a recorded HTTP request
@@ -51,9 +52,9 @@ interface Tape {
 	};
 }
 
-function tapeToResponse(tape: Tape) {
+async function tapeToResponse(tape: Tape) {
 	return new Response(
-		Uint8Array.from(atob(tape.response.body), (c) => c.charCodeAt(0)),
+		tape.response.body?.startsWith(BASE64_PREFIX) ? (await originalFetch(tape.response.body)).body : tape.response.body,
 		{
 			status: tape.response.status,
 			statusText: tape.response.statusText,
@@ -133,21 +134,32 @@ async function vcr(
 	const response = await originalFetch(input, init);
 
 	if (VCR_MODE === MODE.RECORD || VCR_MODE === MODE.CACHE) {
+		const isText =
+			response.headers.get("Content-Type")?.includes("json") || response.headers.get("Content-Type")?.includes("text");
 		const arrayBuffer = await response.arrayBuffer();
-		const headers: Record<string, string> = {};
-		response.headers.forEach((value, key) => (headers[key] = value));
 
 		const tape: Tape = {
 			url,
 			init: {
-				headers: init.headers,
+				headers: omit(init.headers as Record<string, string>, "Authorization"),
 				method: init.method,
 			},
 			response: {
 				// Truncating the body to 30KB to avoid having huge files
-				body: arrayBuffer.byteLength > 30_000 ? "" : Buffer.from(arrayBuffer).toString("base64"),
+				body:
+					arrayBuffer.byteLength > 30_000
+						? ""
+						: isText
+						? new TextDecoder().decode(arrayBuffer)
+						: BASE64_PREFIX + base64FromBytes(new Uint8Array(arrayBuffer)),
 				status: response.status,
 				statusText: response.statusText,
+				headers: Object.fromEntries(
+					// Remove varying headers as much as possible
+					[...response.headers.entries()].filter(
+						([key]) => key !== "date" && key !== "content-length" && !key.startsWith("x-")
+					)
+				),
 			},
 		};
 		tapes[hash] = tape;