fix(ts): extract JSON from chatty LLM responses in fact retrieval (#4533)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>

mem0-ts/src/oss/src/memory/index.ts

@@ -19,7 +19,7 @@ import {
   getFactRetrievalMessages,
   getUpdateMemoryMessages,
   parseMessages,
-  removeCodeBlocks,
+  extractJson,
 } from "../prompts";
 import { DummyHistoryManager } from "../storage/DummyHistoryManager";
 import { Embedder } from "../embeddings/base";
@@ -310,7 +310,7 @@ export class Memory {
       { type: "json_object" },
     );
 
-    const cleanResponse = removeCodeBlocks(response as string);
+    const cleanResponse = extractJson(response as string);
     let facts: string[] = [];
     try {
       const parsed = FactRetrievalSchema.parse(JSON.parse(cleanResponse));
@@ -364,7 +364,7 @@ export class Memory {
       { type: "json_object" },
     );
 
-    const cleanUpdateResponse = removeCodeBlocks(updateResponse as string);
+    const cleanUpdateResponse = extractJson(updateResponse as string);
     let memoryActions: any[] = [];
     try {
       memoryActions = JSON.parse(cleanUpdateResponse).memory || [];

mem0-ts/src/oss/src/prompts/index.ts

@@ -281,5 +281,48 @@ export function parseMessages(messages: string[]): string {
 export function removeCodeBlocks(text: string): string {
   // Extract content inside code fences, handling both complete and
   // truncated blocks (where the closing ``` never arrives).
-  return text.replace(/```(?:\w+)?\n?([\s\S]*?)(?:```|$)/g, "$1").trim();
+  const stripped = text
+    .replace(/```(?:\w+)?\n?([\s\S]*?)(?:```|$)/g, "$1")
+    .trim();
+  // Strip <think>...</think> blocks emitted by reasoning models (e.g. DeepSeek)
+  return stripped.replace(/<think>[\s\S]*?<\/think>/g, "").trim();
+}
+
+/**
+ * Extracts a JSON object from text that may be wrapped in explanation text.
+ *
+ * Some LLMs (especially local models like Ollama/LM Studio) return JSON
+ * wrapped in conversational text without code fences, e.g.:
+ *
+ *   "Here are the facts I extracted:\n{\"facts\": [\"fact1\"]}\nI hope this helps!"
+ *
+ * This function first tries `removeCodeBlocks` for code-fence-wrapped JSON,
+ * then falls back to locating the first `{` and last `}` to extract the
+ * outermost JSON object.
+ *
+ * @param text - The raw LLM response text
+ * @returns The extracted JSON string, or the original text if no JSON object
+ *          boundaries are found
+ */
+export function extractJson(text: string): string {
+  // Step 1: Strip code fences if present
+  const cleaned = removeCodeBlocks(text);
+  const trimmed = cleaned.trim();
+
+  // Step 2: Try to locate a JSON object by first `{` and last `}` boundaries
+  const firstBrace = trimmed.indexOf("{");
+  const lastBrace = trimmed.lastIndexOf("}");
+  if (firstBrace !== -1 && lastBrace > firstBrace) {
+    return trimmed.substring(firstBrace, lastBrace + 1);
+  }
+
+  // Step 3: Try to locate a JSON array by first `[` and last `]` boundaries
+  const firstBracket = trimmed.indexOf("[");
+  const lastBracket = trimmed.lastIndexOf("]");
+  if (firstBracket !== -1 && lastBracket > firstBracket) {
+    return trimmed.substring(firstBracket, lastBracket + 1);
+  }
+
+  // No JSON boundaries found — return as-is and let the caller handle the error
+  return trimmed;
 }

mem0-ts/src/oss/tests/extract-json.test.ts

@@ -0,0 +1,151 @@
+import { extractJson } from "../src/prompts";
+
+describe("extractJson", () => {
+  it("returns clean JSON unchanged", () => {
+    const input = '{"facts": ["hello", "world"]}';
+    expect(extractJson(input)).toBe('{"facts": ["hello", "world"]}');
+  });
+
+  it("extracts JSON from ```json code fence", () => {
+    const input = '```json\n{"facts": ["hello"]}\n```';
+    expect(extractJson(input)).toBe('{"facts": ["hello"]}');
+  });
+
+  it("extracts JSON from bare ``` code fence", () => {
+    const input = '```\n{"facts": ["test"]}\n```';
+    expect(extractJson(input)).toBe('{"facts": ["test"]}');
+  });
+
+  it("extracts JSON wrapped in explanation text without code fences", () => {
+    const input =
+      'Here are the facts I extracted:\n{"facts": ["fact1", "fact2"]}\nI hope this helps!';
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({ facts: ["fact1", "fact2"] });
+  });
+
+  it("extracts JSON from chatty LLM response with leading text", () => {
+    const input =
+      'Based on the conversation, here is the JSON output:\n{"facts": ["Name is John", "Is a software engineer"]}';
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({
+      facts: ["Name is John", "Is a software engineer"],
+    });
+  });
+
+  it("extracts JSON from chatty LLM response with trailing text", () => {
+    const input =
+      '{"facts": ["Loves pizza"]}\nLet me know if you need anything else!';
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({ facts: ["Loves pizza"] });
+  });
+
+  it("extracts JSON from text with both leading and trailing explanation", () => {
+    const input =
+      "Sure! Here's the extracted information:\n" +
+      '{"memory": [{"id": "0", "text": "Name is John", "event": "NONE"}]}\n' +
+      "I've analyzed the conversation above.";
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({
+      memory: [{ id: "0", text: "Name is John", event: "NONE" }],
+    });
+  });
+
+  it("extracts JSON from code-fenced response with surrounding text", () => {
+    const input =
+      'Here is the JSON:\n```json\n{"facts": ["user likes TypeScript"]}\n```\nHope this helps!';
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({
+      facts: ["user likes TypeScript"],
+    });
+  });
+
+  it("handles nested JSON objects", () => {
+    const input =
+      'The output is: {"memory": [{"id": "0", "text": "test", "event": "ADD"}]}';
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({
+      memory: [{ id: "0", text: "test", event: "ADD" }],
+    });
+  });
+
+  it("handles multi-line JSON in chatty text", () => {
+    const input = `Here are the facts:
+{
+  "facts": [
+    "Sky is blue",
+    "Grass is green"
+  ]
+}
+That's all I found.`;
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({
+      facts: ["Sky is blue", "Grass is green"],
+    });
+  });
+
+  it("returns original text when no JSON boundaries found", () => {
+    const input = "No JSON here, just some plain text response.";
+    expect(extractJson(input)).toBe(
+      "No JSON here, just some plain text response.",
+    );
+  });
+
+  it("handles JSON array responses", () => {
+    const input = 'The results are: ["fact1", "fact2", "fact3"]';
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual(["fact1", "fact2", "fact3"]);
+  });
+
+  it("returns empty string for empty input", () => {
+    expect(extractJson("")).toBe("");
+  });
+
+  it("handles truncated code block missing closing fence", () => {
+    const input = '```json\n{"facts": ["hello"]}';
+    expect(extractJson(input)).toBe('{"facts": ["hello"]}');
+  });
+
+  it("handles whitespace-padded JSON", () => {
+    const input = '   {"facts": ["test"]}   ';
+    expect(extractJson(input)).toBe('{"facts": ["test"]}');
+  });
+
+  it("handles LM Studio-style verbose response", () => {
+    const input =
+      "I'll analyze the conversation and extract the relevant facts.\n\n" +
+      '{"facts": ["User prefers dark mode", "User uses VS Code"]}\n\n' +
+      "These are the key preferences I identified from the conversation.";
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({
+      facts: ["User prefers dark mode", "User uses VS Code"],
+    });
+  });
+
+  it("handles Ollama-style response with thinking prefix", () => {
+    const input =
+      "Let me think about this...\n\n" +
+      "After analyzing the input, here is my response:\n" +
+      '{"facts": ["Has a dog named Max"]}';
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({
+      facts: ["Has a dog named Max"],
+    });
+  });
+
+  it("strips <think> blocks from reasoning models before extracting JSON", () => {
+    const input =
+      "<think>\nLet me analyze the conversation carefully.\n</think>\n" +
+      '{"facts": ["User lives in Tokyo"]}';
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({
+      facts: ["User lives in Tokyo"],
+    });
+  });
+
+  it("handles <think> blocks inside code fences", () => {
+    const input =
+      '```json\n<think>reasoning here</think>\n{"facts": ["test"]}\n```';
+    const result = extractJson(input);
+    expect(JSON.parse(result)).toEqual({ facts: ["test"] });
+  });
+});

mem0-ts/src/oss/tests/remove-code-blocks.test.ts

@@ -65,4 +65,23 @@ describe("removeCodeBlocks", () => {
     const input = '```json\r\n{"facts": ["hello"]}\r\n```';
     expect(removeCodeBlocks(input)).toBe('{"facts": ["hello"]}');
   });
+
+  it("strips <think> blocks from reasoning models", () => {
+    const input =
+      '<think>Let me analyze this conversation...</think>\n{"facts": ["Name is John"]}';
+    expect(removeCodeBlocks(input)).toBe('{"facts": ["Name is John"]}');
+  });
+
+  it("strips <think> blocks inside code fences", () => {
+    const input =
+      '```json\n<think>thinking about it</think>\n{"facts": ["test"]}\n```';
+    expect(removeCodeBlocks(input)).toBe('{"facts": ["test"]}');
+  });
+
+  it("strips multi-line <think> blocks", () => {
+    const input =
+      "<think>\nStep 1: Read the conversation\nStep 2: Extract facts\n</think>\n" +
+      '{"facts": ["Likes pizza"]}';
+    expect(removeCodeBlocks(input)).toBe('{"facts": ["Likes pizza"]}');
+  });
 });