fix(oss): validate LLM fact output via FactRetrievalSchema before embedding (#4083)

mgoulart · web-flow · commit db15d5c6297e · 2026-02-22T19:57:48.000-08:00
diff --git a/mem0-ts/src/oss/src/embeddings/ollama.ts b/mem0-ts/src/oss/src/embeddings/ollama.ts
@@ -27,9 +27,12 @@ export class OllamaEmbedder implements Embedder {
     } catch (err) {
       logger.error(`Error ensuring model exists: ${err}`);
     }
+    // Ollama's Go server requires prompt to be a string. Coerce defensively
+    // since callers may pass values parsed from untrusted LLM JSON output.
+    const prompt = typeof text === "string" ? text : JSON.stringify(text);
     const response = await this.ollama.embeddings({
       model: this.model,
-      prompt: text,
+      prompt,
     });
     return response.embedding;
   }
diff --git a/mem0-ts/src/oss/src/memory/index.ts b/mem0-ts/src/oss/src/memory/index.ts
@@ -15,6 +15,7 @@ import {
   HistoryManagerFactory,
 } from "../utils/factory";
 import {
+  FactRetrievalSchema,
   getFactRetrievalMessages,
   getUpdateMemoryMessages,
   parseMessages,
@@ -261,7 +262,8 @@ export class Memory {
     const cleanResponse = removeCodeBlocks(response as string);
     let facts: string[] = [];
     try {
-      facts = JSON.parse(cleanResponse).facts || [];
+      const parsed = FactRetrievalSchema.parse(JSON.parse(cleanResponse));
+      facts = parsed.facts;
     } catch (e) {
       console.error(
         "Failed to parse facts from LLM response:",
diff --git a/mem0-ts/src/oss/src/prompts/index.ts b/mem0-ts/src/oss/src/prompts/index.ts
@@ -1,9 +1,18 @@
 import { z } from "zod";
 
+// Accepts a string directly, or an object with a "fact" or "text" key
+// (common malformed shapes from smaller LLMs like llama3.1:8b).
+const factItem = z.union([
+  z.string(),
+  z.object({ fact: z.string() }).transform((o) => o.fact),
+  z.object({ text: z.string() }).transform((o) => o.text),
+]);
+
 // Define Zod schema for fact retrieval output
 export const FactRetrievalSchema = z.object({
   facts: z
-    .array(z.string())
+    .array(factItem)
+    .transform((arr) => arr.filter((s) => s.length > 0))
     .describe("An array of distinct facts extracted from the conversation."),
 });
 

Original file line number	Diff line number	Diff line change
`@@ -27,9 +27,12 @@ export class OllamaEmbedder implements Embedder {`
`27`	`27`	`} catch (err) {`
`28`	`28`	logger.error(`Error ensuring model exists: ${err}`);
`29`	`29`	`}`
	`30`	`+ // Ollama's Go server requires prompt to be a string. Coerce defensively`
	`31`	`+ // since callers may pass values parsed from untrusted LLM JSON output.`
	`32`	`+ const prompt = typeof text === "string" ? text : JSON.stringify(text);`
`30`	`33`	`const response = await this.ollama.embeddings({`
`31`	`34`	`model: this.model,`
`32`		`- prompt: text,`
	`35`	`+ prompt,`
`33`	`36`	`});`
`34`	`37`	`return response.embedding;`
`35`	`38`	`}`