diff --git a/sdk/wren-langchain/src/wren_langchain/_format.py b/sdk/wren-langchain/src/wren_langchain/_format.py index 82d3085197..223c2db576 100644 --- a/sdk/wren-langchain/src/wren_langchain/_format.py +++ b/sdk/wren-langchain/src/wren_langchain/_format.py @@ -110,10 +110,19 @@ def format_recall_content(rows: list[dict[str, Any]]) -> str: return "_No similar past queries found._" chunks = [] - for i, row in enumerate(rows, start=1): + n = 0 + for row in rows: + # Vector stores / serializers occasionally emit None or plain + # strings alongside dict hits. Calling ``.get`` on those aborted + # the entire recall content render. + if not isinstance(row, dict): + continue nl = row.get("nl_query") or row.get("nl") or "" sql = row.get("sql_query") or row.get("sql") or "" - chunks.append(f'{i}. "{nl}"\n ```sql\n {sql}\n ```') + n += 1 + chunks.append(f'{n}. "{nl}"\n ```sql\n {sql}\n ```') + if not chunks: + return "_No similar past queries found._" return "\n".join(chunks) diff --git a/sdk/wren-langchain/tests/unit/test_format_recall_content_guard.py b/sdk/wren-langchain/tests/unit/test_format_recall_content_guard.py new file mode 100644 index 0000000000..1b796dea5e --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_format_recall_content_guard.py @@ -0,0 +1,26 @@ +import importlib.util +from pathlib import Path + +_ROOT = Path(__file__).resolve().parents[2] +_PATH = _ROOT / "src" / "wren_langchain" / "_format.py" +_spec = importlib.util.spec_from_file_location("_fmt", _PATH) +_mod = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_mod) + + +def test_skips_non_dict_rows(): + out = _mod.format_recall_content( + [ + None, + "x", + {"nl": "q1", "sql": "SELECT 1"}, + {"nl_query": "q2", "sql_query": "SELECT 2"}, + ] + ) + assert "q1" in out and "SELECT 1" in out + assert "q2" in out + assert out.startswith('1. "') + + +def test_all_invalid_returns_empty_message(): + assert _mod.format_recall_content([None, 3]) == "_No similar past queries found._"