fix: handle chatty LLM responses in JSON parsing (#4525)
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Utkarsh
+18
-12
@@ -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 = {}
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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\n<think>reasoning here</think>\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 = '<think>reasoning here</think>\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 <think> tags."""
|
||||
response = '<think>Let me process this...</think>\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"]
|
||||
Reference in New Issue
Block a user