From f1b337a975949ecac56d83c9940807653b79548b Mon Sep 17 00:00:00 2001
From: santosh2345 <santoshgiri2345@gmail.com>
Date: Thu, 16 Jul 2026 12:21:00 +0545
Subject: [PATCH] fix(extractor): don't skip JSON responses that carry natural
 language

_is_trivial_response() blanket-skipped any response starting with '{' or
'[', so an LLM answer wrapped in JSON like
{"answer": "the user's budget is $20k, they prefer Python"} was
classified as trivial and never extracted. Closes #24.

- Replace the naive startswith('{'/'[') check with classify_response_type:
  skip pure "sql"/"code" payloads, but for "json" only skip when the
  payload's string values contain no meaningful prose (< 5 word-like
  tokens). Pure data payloads such as {"key": "value"} stay trivial.
- Make the code-character density heuristic whitespace-insensitive so
  symbol soup like "{ } ( ) = ;" is still caught now that the startswith
  guard is gone.
- Tests: JSON-with-prose is non-trivial (regression), pure-data JSON stays
  trivial, SQL and code fences stay trivial.
---
 beliefstate/extractor.py   | 46 ++++++++++++++++++++++++++++++++++----
 tests/test_new_features.py | 23 +++++++++++++++++++
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/beliefstate/extractor.py b/beliefstate/extractor.py
index 6a17852..4e60024 100644
--- a/beliefstate/extractor.py
+++ b/beliefstate/extractor.py
@@ -154,6 +154,37 @@ def classify_response_type(text: str) -> str:
     return "conversational"
 
 
+_MIN_JSON_PROSE_WORDS = 5
+
+
+def _json_prose_word_count(text: str) -> int:
+    """Count word-like tokens inside a JSON payload's string *values*.
+
+    Object keys are treated as structural and ignored; only values are walked.
+    Returns 0 if the text is not valid JSON.
+    """
+    try:
+        data = json.loads(text)
+    except (json.JSONDecodeError, ValueError):
+        return 0
+
+    words = 0
+
+    def walk(obj: Any) -> None:
+        nonlocal words
+        if isinstance(obj, str):
+            words += len(re.findall(r"[A-Za-z]{2,}", obj))
+        elif isinstance(obj, dict):
+            for value in obj.values():
+                walk(value)
+        elif isinstance(obj, list):
+            for value in obj:
+                walk(value)
+
+    walk(data)
+    return words
+
+
 def _is_trivial_response(text: str) -> bool:
     """Check if an assistant response is trivial and should be skipped for extraction."""
     if not text or not text.strip():
@@ -168,12 +199,19 @@ def _is_trivial_response(text: str) -> bool:
             if re.match(pattern, text_lower):
                 return True
     code_chars = set("{}()=;:|#|\\><")
-    if text:
-        code_count = sum(1 for c in text if c in code_chars)
-        if code_count / len(text) > 0.6:
+    non_space = [c for c in text if not c.isspace()]
+    if non_space:
+        code_count = sum(1 for c in non_space if c in code_chars)
+        if code_count / len(non_space) > 0.6:
             return True
-    if text.startswith("{") or text.startswith("["):
+    # A structurally JSON/SQL/code response is not automatically trivial: an LLM
+    # may wrap real answers in JSON (e.g. {"answer": "the budget is $20k ..."}).
+    # Skip pure data payloads and code output, but keep JSON that carries prose.
+    response_type = classify_response_type(text)
+    if response_type in ("sql", "code"):
         return True
+    if response_type == "json":
+        return _json_prose_word_count(text) < _MIN_JSON_PROSE_WORDS
     return False
 
 
diff --git a/tests/test_new_features.py b/tests/test_new_features.py
index d231848..2ec0bc4 100644
--- a/tests/test_new_features.py
+++ b/tests/test_new_features.py
@@ -134,9 +134,32 @@ def test_is_trivial_code_block(self):
         assert _is_trivial_response("{ } ( ) = ; : # | \\ > <") is True
 
     def test_is_trivial_json_start(self):
+        # Pure data payloads (no meaningful prose) stay trivial.
         assert _is_trivial_response('{"key": "value"}') is True
         assert _is_trivial_response('[{"key": "value"}]') is True
 
+    def test_not_trivial_json_with_natural_language(self):
+        # Regression for #24: an LLM answer wrapped in JSON must NOT be skipped.
+        assert (
+            _is_trivial_response(
+                '{"answer": "The user\'s budget is $20k, they prefer Python"}'
+            )
+            is False
+        )
+        assert (
+            _is_trivial_response(
+                '{"reply": "We decided to use PostgreSQL for the main database"}'
+            )
+            is False
+        )
+
+    def test_is_trivial_sql(self):
+        assert _is_trivial_response("SELECT id, name FROM users WHERE id = 1") is True
+
+    def test_is_trivial_code_fence(self):
+        code = "```python\n" + "print('hello world')\n" * 3 + "```"
+        assert _is_trivial_response(code) is True
+
     def test_not_trivial_substantive(self):
         assert _is_trivial_response("I understand your budget constraints.") is False
         assert _is_trivial_response("The database should be PostgreSQL.") is False