Browse Source

Fix problems in JSON parsing.

saul 1 year ago
parent
commit
825e70954f
3 changed files with 56 additions and 40 deletions
  1. 1 1
      package.json
  2. 27 39
      src/index.ts
  3. 28 0
      src/utils.ts

+ 1 - 1
package.json

@@ -1,7 +1,7 @@
 {
 {
   "type": "module",
   "type": "module",
   "name": "ollama",
   "name": "ollama",
-  "version": "0.2.0",
+  "version": "0.2.1",
   "description": "Interface with an ollama instance over HTTP.",
   "description": "Interface with an ollama instance over HTTP.",
   "main": "dist/index.js",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "types": "dist/index.d.ts",

+ 27 - 39
src/index.ts

@@ -54,27 +54,23 @@ export class Ollama {
 			throw new Error("Missing body");
 			throw new Error("Missing body");
 		}
 		}
 
 
-		for await (const chunk of response.body) {
-			const messages = chunk.toString().split("\n").filter(s => s.length !== 0);
-
-			for (const message of messages) {
-				const res: GenerateResponse | GenerateResponseEnd = JSON.parse(message);
-
-				if (res.done) {
-					return {
-						model: res.model,
-						createdAt: new Date(res.created_at),
-						context: res.context,
-						totalDuration: res.total_duration,
-						loadDuration: res.load_duration,
-						promptEvalCount: res.prompt_eval_count,
-						evalCount: res.eval_count,
-						evalDuration: res.eval_duration
-					};
-				}
-
-				yield res.response;
+		const itr = utils.parseJSON<GenerateResponse | GenerateResponseEnd>(response.body);
+
+		for await (const message of itr) {
+			if (message.done) {
+				return {
+					model: message.model,
+					createdAt: new Date(message.created_at),
+					context: message.context,
+					totalDuration: message.total_duration,
+					loadDuration: message.load_duration,
+					promptEvalCount: message.prompt_eval_count,
+					evalCount: message.eval_count,
+					evalDuration: message.eval_duration
+				};
 			}
 			}
+
+			yield message.response;
 		}
 		}
 
 
 		throw new Error("Did not recieve done response in stream.");
 		throw new Error("Did not recieve done response in stream.");
@@ -87,14 +83,10 @@ export class Ollama {
 			throw new Error("Missing body");
 			throw new Error("Missing body");
 		}
 		}
 
 
-		for await (const chunk of response.body) {
-			const messages = chunk.toString().split("\n").filter(s => s.length !== 0);
+		const itr = utils.parseJSON<CreateResponse>(response.body);
 
 
-			for (const message of messages) {
-				const res: CreateResponse = JSON.parse(message);
-
-				yield res.status;
-			}
+		for await (const message of itr) {
+			yield message.status;
 		}
 		}
 	}
 	}
 
 
@@ -116,19 +108,15 @@ export class Ollama {
 			throw new Error("Missing body");
 			throw new Error("Missing body");
 		}
 		}
 
 
-		for await (const chunk of response.body) {
-			const messages = chunk.toString().split("\n").filter(s => s.length !== 0);
+		const itr = utils.parseJSON<PullResponse>(response.body);
 
 
-			for (const message of messages) {
-				const res: PullResponse = JSON.parse(message);
-
-				yield {
-					status: res.status,
-					digest: res["digest"] ?? "",
-					total: res["total"] ?? 0,
-					completed: res["completed"] ?? 0
-				};
-			}
+		for await (const message of itr) {
+			yield {
+				status: message.status,
+				digest: message["digest"] ?? "",
+				total: message["total"] ?? 0,
+				completed: message["completed"] ?? 0
+			};
 		}
 		}
 	}
 	}
 
 

+ 28 - 0
src/utils.ts

@@ -57,3 +57,31 @@ export const del = async (address: string, data?: Record<string, unknown>): Prom
 
 
 	return response;
 	return response;
 };
 };
+
+export const parseJSON = async function * <T = unknown>(itr: Iterable<{ toString: () => string }> | AsyncIterable<{ toString: () => string }>): AsyncGenerator<T> {
+	let buffer = "";
+
+	for await (const chunk of itr) {
+		buffer += chunk;
+
+		const parts = buffer.split("\n");
+
+		buffer = parts.pop() ?? "";
+
+		for (const part of parts) {
+			try {
+				yield JSON.parse(part);
+			} catch (error) {
+				console.warn("invalid json: ", part);
+			}
+		}
+	}
+
+	for (const part of buffer.split("\n").filter(p => p !== "")) {
+		try {
+			yield JSON.parse(part);
+		} catch (error) {
+			console.warn("invalid json: ", part);
+		}
+	}
+};