From 2a59c9fd99b4a69d83841e575542d42a748953d4 Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Thu, 26 Mar 2026 16:39:17 +0530 Subject: [PATCH] fix: handle chatty LLM responses in JSON parsing (#4525) Co-authored-by: utkarsh240799 Co-authored-by: Claude Opus 4.6 (1M context) --- mem0/memory/main.py | 30 ++-- mem0/memory/utils.py | 10 +- tests/test_chatty_llm_parsing.py | 232 +++++++++++++++++++++++++++++++ 3 files changed, 258 insertions(+), 14 deletions(-) create mode 100644 tests/test_chatty_llm_parsing.py diff --git a/mem0/memory/main.py b/mem0/memory/main.py index 31d638223..78e50fb39 100644 --- a/mem0/memory/main.py +++ b/mem0/memory/main.py @@ -516,15 +516,15 @@ class Memory(MemoryBase): ) try: - response = remove_code_blocks(response) - if not response.strip(): + cleaned_response = remove_code_blocks(response) + if not cleaned_response.strip(): new_retrieved_facts = [] else: try: # First try direct JSON parsing - new_retrieved_facts = json.loads(response, strict=False)["facts"] + new_retrieved_facts = json.loads(cleaned_response, strict=False)["facts"] except json.JSONDecodeError: - # Try extracting JSON from response using built-in function + # Try extracting JSON from response (handles chatty LLM output) extracted_json = extract_json(response) new_retrieved_facts = json.loads(extracted_json, strict=False)["facts"] new_retrieved_facts = normalize_facts(new_retrieved_facts) @@ -589,8 +589,11 @@ class Memory(MemoryBase): logger.warning("Empty response from LLM, no memories to extract") new_memories_with_actions = {} else: - response = remove_code_blocks(response) - new_memories_with_actions = json.loads(response, strict=False) + try: + new_memories_with_actions = json.loads(remove_code_blocks(response), strict=False) + except json.JSONDecodeError: + extracted_json = extract_json(response) + new_memories_with_actions = json.loads(extracted_json, strict=False) except Exception as e: logger.error(f"Invalid JSON response: {e}") new_memories_with_actions = {} @@ -1590,15 +1593,15 @@ class AsyncMemory(MemoryBase): response_format={"type": "json_object"}, ) try: - response = remove_code_blocks(response) - if not response.strip(): + cleaned_response = remove_code_blocks(response) + if not cleaned_response.strip(): new_retrieved_facts = [] else: try: # First try direct JSON parsing - new_retrieved_facts = json.loads(response, strict=False)["facts"] + new_retrieved_facts = json.loads(cleaned_response, strict=False)["facts"] except json.JSONDecodeError: - # Try extracting JSON from response using built-in function + # Try extracting JSON from response (handles chatty LLM output) extracted_json = extract_json(response) new_retrieved_facts = json.loads(extracted_json, strict=False)["facts"] new_retrieved_facts = normalize_facts(new_retrieved_facts) @@ -1666,8 +1669,11 @@ class AsyncMemory(MemoryBase): logger.warning("Empty response from LLM, no memories to extract") new_memories_with_actions = {} else: - response = remove_code_blocks(response) - new_memories_with_actions = json.loads(response, strict=False) + try: + new_memories_with_actions = json.loads(remove_code_blocks(response), strict=False) + except json.JSONDecodeError: + extracted_json = extract_json(response) + new_memories_with_actions = json.loads(extracted_json, strict=False) except Exception as e: logger.error(f"Invalid JSON response: {e}") new_memories_with_actions = {} diff --git a/mem0/memory/utils.py b/mem0/memory/utils.py index e0581a09c..3a353144a 100644 --- a/mem0/memory/utils.py +++ b/mem0/memory/utils.py @@ -124,14 +124,20 @@ def remove_code_blocks(content: str) -> str: def extract_json(text): """ Extracts JSON content from a string, removing enclosing triple backticks and optional 'json' tag if present. - If no code block is found, returns the text as-is. + If no code block is found, attempts to locate JSON by finding the first '{' and last '}'. + If that also fails, returns the text as-is. """ text = text.strip() match = re.search(r"```(?:json)?\s*(.*?)\s*```", text, re.DOTALL) if match: json_str = match.group(1) else: - json_str = text # assume it's raw JSON + start_idx = text.find("{") + end_idx = text.rfind("}") + if start_idx != -1 and end_idx != -1 and end_idx > start_idx: + json_str = text[start_idx : end_idx + 1] + else: + json_str = text return json_str diff --git a/tests/test_chatty_llm_parsing.py b/tests/test_chatty_llm_parsing.py new file mode 100644 index 000000000..38d4c061a --- /dev/null +++ b/tests/test_chatty_llm_parsing.py @@ -0,0 +1,232 @@ +"""Tests for chatty LLM JSON parsing robustness. + +Validates that extract_json and the full fallback chain correctly handle +responses from local LLMs (LM Studio, Ollama) that wrap JSON in +conversational text, markdown code blocks, or both. +""" + +import json + +import pytest + +from mem0.memory.utils import extract_json, remove_code_blocks + + +# --- Test extract_json --- + +class TestExtractJson: + """Tests for extract_json utility.""" + + def test_pure_json(self): + text = '{"memory": [{"id": "0", "text": "likes basketball", "event": "ADD"}]}' + result = extract_json(text) + parsed = json.loads(result) + assert parsed["memory"][0]["text"] == "likes basketball" + + def test_json_in_markdown_code_block(self): + text = '```json\n{"memory": [{"id": "0", "text": "likes basketball", "event": "ADD"}]}\n```' + result = extract_json(text) + parsed = json.loads(result) + assert parsed["memory"][0]["text"] == "likes basketball" + + def test_json_in_plain_code_block(self): + text = '```\n{"memory": [{"id": "0", "text": "likes basketball", "event": "ADD"}]}\n```' + result = extract_json(text) + parsed = json.loads(result) + assert parsed["memory"][0]["text"] == "likes basketball" + + def test_chatty_with_markdown(self): + """extract_json can handle embedded code blocks via re.search.""" + text = """Here is the extracted memory: +```json +{"memory": [{"id": "0", "text": "likes basketball", "event": "ADD"}]} +``` +I hope this helps!""" + result = extract_json(text) + parsed = json.loads(result) + assert parsed["memory"][0]["text"] == "likes basketball" + + def test_chatty_without_markdown(self): + """extract_json finds JSON by {/} boundaries when no code block is present.""" + text = """Here is the memory update you requested: +{"memory": [{"id": "0", "text": "likes gaming", "event": "ADD"}]} +That's the result.""" + result = extract_json(text) + parsed = json.loads(result) + assert parsed["memory"][0]["text"] == "likes gaming" + + def test_chatty_multiline_json_without_markdown(self): + """extract_json finds multi-line JSON by {/} boundaries.""" + text = """Here is the memory update: +{ + "memory": [ + { + "id": "0", + "text": "User likes gaming", + "event": "ADD" + } + ] +} +""" + result = extract_json(text) + parsed = json.loads(result) + assert parsed["memory"][0]["text"] == "User likes gaming" + + def test_no_json_at_all(self): + """No JSON in the text — should return as-is.""" + text = "I don't have any memory updates." + result = extract_json(text) + assert result == text + + def test_whitespace_padding(self): + text = ' \n {"memory": []} \n ' + result = extract_json(text) + parsed = json.loads(result) + assert parsed == {"memory": []} + + def test_nested_json_objects(self): + """JSON with nested objects should find outermost braces.""" + text = 'Sure! {"memory": [{"id": "0", "text": "test", "event": "ADD"}]} Done.' + result = extract_json(text) + parsed = json.loads(result) + assert parsed["memory"][0]["id"] == "0" + + +# --- Test remove_code_blocks --- + +class TestRemoveCodeBlocks: + """Tests for remove_code_blocks — verify it does NOT handle chatty text.""" + + def test_clean_code_block(self): + text = '```json\n{"memory": []}\n```' + result = remove_code_blocks(text) + parsed = json.loads(result) + assert parsed == {"memory": []} + + def test_no_code_block(self): + """Without code blocks, returns content as-is (may not be valid JSON).""" + text = 'Here is the result: {"memory": []}' + result = remove_code_blocks(text) + # This should NOT be parseable as JSON because it has leading text + with pytest.raises(json.JSONDecodeError): + json.loads(result) + + def test_think_tags_removed_inside_code_block(self): + """Think tags inside a code block response are removed.""" + text = '```json\nreasoning here\n{"memory": []}\n```' + result = remove_code_blocks(text) + parsed = json.loads(result) + assert parsed == {"memory": []} + + def test_think_tags_before_code_block_not_handled(self): + """Think tags before code block cause regex to miss — fallback needed. + + This is a known limitation of remove_code_blocks: it matches the code + block regex first, and only strips think tags from the result. When + think tags come before the code block, the regex fails entirely. + The fallback chain (extract_json) handles this case. + """ + text = 'reasoning here\n```json\n{"memory": []}\n```' + result = remove_code_blocks(text) + with pytest.raises(json.JSONDecodeError): + json.loads(result) + + +# --- Test the full fallback chain (remove_code_blocks -> extract_json) --- + +class TestFallbackChain: + """Tests the actual fallback pattern used in _add_to_vector_store: + try json.loads(remove_code_blocks(response)) + except JSONDecodeError: json.loads(extract_json(response)) + """ + + def _parse_with_fallback(self, response): + """Mimics the parsing logic in _add_to_vector_store.""" + try: + return json.loads(remove_code_blocks(response), strict=False) + except json.JSONDecodeError: + extracted = extract_json(response) + return json.loads(extracted, strict=False) + + def test_clean_json(self): + response = '{"memory": [{"id": "0", "text": "Name is Alex", "event": "ADD"}]}' + result = self._parse_with_fallback(response) + assert result["memory"][0]["text"] == "Name is Alex" + + def test_markdown_wrapped(self): + response = '```json\n{"memory": [{"id": "0", "text": "Name is Alex", "event": "ADD"}]}\n```' + result = self._parse_with_fallback(response) + assert result["memory"][0]["text"] == "Name is Alex" + + def test_chatty_markdown(self): + """Issue #3788: LLM wraps in conversational text + markdown.""" + response = """Here is the extracted memory: +```json +{ + "memory": [ + {"id": "0", "text": "Name is Alex", "event": "ADD"}, + {"id": "1", "text": "Love basketball", "event": "ADD"}, + {"id": "2", "text": "Love gaming", "event": "ADD"} + ] +} +``` +I hope this helps!""" + result = self._parse_with_fallback(response) + assert len(result["memory"]) == 3 + assert result["memory"][0]["text"] == "Name is Alex" + assert result["memory"][1]["text"] == "Love basketball" + assert result["memory"][2]["text"] == "Love gaming" + + def test_chatty_no_markdown(self): + """Issue #3788: LLM wraps in conversational text without markdown.""" + response = """Here is the memory update you requested: +{ + "memory": [ + {"id": "0", "text": "Name is Alex", "event": "ADD"}, + {"id": "1", "text": "Love basketball", "event": "ADD"} + ] +}""" + result = self._parse_with_fallback(response) + assert len(result["memory"]) == 2 + + def test_think_tags_with_json(self): + """Reasoning model with tags.""" + response = 'Let me process this...\n```json\n{"memory": [{"id": "0", "text": "test", "event": "ADD"}]}\n```' + result = self._parse_with_fallback(response) + assert result["memory"][0]["text"] == "test" + + def test_lmstudio_style_response(self): + """Mimics the actual LM Studio response format from the issue.""" + response = """{ + "memory": [ + { + "id": "0", + "text": "Name is Alex", + "event": "ADD" + }, + { + "id": "1", + "text": "Love basketball", + "event": "ADD" + }, + { + "id": "2", + "text": "Love gaming", + "event": "ADD" + } + ] +}""" + result = self._parse_with_fallback(response) + assert len(result["memory"]) == 3 + + def test_completely_invalid_response(self): + """If both attempts fail, the outer except in main.py catches it.""" + response = "I don't understand the question" + with pytest.raises(json.JSONDecodeError): + self._parse_with_fallback(response) + + def test_facts_extraction_fallback(self): + """Also verify the pattern works for facts extraction responses.""" + response = 'Sure! Here are the facts:\n{"facts": ["Name is Alex", "Loves basketball"]}\nHope that helps!' + result = self._parse_with_fallback(response) + assert result["facts"] == ["Name is Alex", "Loves basketball"]