diff --git a/sdk/wren-langchain/src/wren_langchain/_format.py b/sdk/wren-langchain/src/wren_langchain/_format.py index 82d3085197..7398a0ccbc 100644 --- a/sdk/wren-langchain/src/wren_langchain/_format.py +++ b/sdk/wren-langchain/src/wren_langchain/_format.py @@ -75,17 +75,25 @@ def format_fetch_context_content(result: dict[str, Any]) -> str: return _cap_to_bytes(text, suffix="\n\n...[truncated]") items = result.get("results", []) or [] - if not items: + if not isinstance(items, list) or not items: return "_No relevant context items found._" lines = [] - for i, item in enumerate(items, start=1): + n = 0 + for item in items: + if not isinstance(item, dict): + continue + n += 1 item_type = item.get("item_type", "item") name = item.get("name", "") summary = item.get("summary") or item.get("text") or "" + if not isinstance(summary, str): + summary = str(summary) if summary is not None else "" if len(summary) > 120: summary = summary[:117] + "..." - lines.append(f"{i}. [{item_type}] {name} — {summary}") + lines.append(f"{n}. [{item_type}] {name} — {summary}") + if not lines: + return "_No relevant context items found._" return _cap_to_bytes("\n".join(lines), suffix="\n...[truncated]") @@ -106,14 +114,20 @@ def _cap_to_bytes(text: str, *, suffix: str) -> str: def format_recall_content(rows: list[dict[str, Any]]) -> str: """Render recalled NL→SQL pairs as a numbered list with code fences.""" - if not rows: + if not isinstance(rows, list) or not rows: return "_No similar past queries found._" chunks = [] - for i, row in enumerate(rows, start=1): + n = 0 + for row in rows: + if not isinstance(row, dict): + continue + n += 1 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 ```') + chunks.append(f'{n}. "{nl}"\n ```sql\n {sql}\n ```') + if not chunks: + return "_No similar past queries found._" return "\n".join(chunks) @@ -132,18 +146,28 @@ def format_list_models_content(manifest: dict[str, Any]) -> str: Columns: model | cols | description. """ models = manifest.get("models", []) or [] - if not models: + if not isinstance(models, list) or not models: return "_No models defined in this Wren project._" lines = ["| model | cols | description |", "|---|---|---|"] + any_model = False for m in models: + if not isinstance(m, dict): + continue + any_model = True name = m.get("name", "") - col_count = len(m.get("columns", []) or []) - desc = ( - (m.get("properties") or {}).get("description") or m.get("description") or "" - ) + cols = m.get("columns", []) or [] + col_count = len(cols) if isinstance(cols, list) else 0 + props = m.get("properties") or {} + if not isinstance(props, dict): + props = {} + desc = props.get("description") or m.get("description") or "" + if not isinstance(desc, str): + desc = str(desc) if desc is not None else "" # Trim long descriptions to keep table compact. if len(desc) > 80: desc = desc[:77] + "..." lines.append(f"| {name} | {col_count} | {desc} |") + if not any_model: + return "_No models defined in this Wren project._" return "\n".join(lines) diff --git a/sdk/wren-langchain/tests/unit/test_format_malformed.py b/sdk/wren-langchain/tests/unit/test_format_malformed.py new file mode 100644 index 0000000000..b5b88c32b9 --- /dev/null +++ b/sdk/wren-langchain/tests/unit/test_format_malformed.py @@ -0,0 +1,123 @@ +"""format_* helpers must tolerate non-dict list rows from memory/search APIs.""" + +from __future__ import annotations + +from wren_langchain._format import ( + format_fetch_context_content, + format_list_models_content, + format_recall_content, +) + +_FETCH_FALLBACK = "_No relevant context items found._" +_RECALL_FALLBACK = "_No similar past queries found._" +_MODELS_FALLBACK = "_No models defined in this Wren project._" + + +def test_format_fetch_context_skips_non_dict_items() -> None: + out = format_fetch_context_content( + { + "strategy": "search", + "results": [ + "not-a-dict", + None, + { + "item_type": "model", + "name": "orders", + "summary": "Orders table", + }, + ], + } + ) + # Skipped rows must not shift the numbering — the first valid row is `1.`. + assert "1. [model] orders" in out + assert "not-a-dict" not in out + + +def test_format_fetch_context_non_list_container_falls_back() -> None: + out = format_fetch_context_content({"strategy": "search", "results": "oops"}) + assert out == _FETCH_FALLBACK + + +def test_format_fetch_context_all_invalid_falls_back() -> None: + out = format_fetch_context_content( + {"strategy": "search", "results": ["x", None, 3]} + ) + assert out == _FETCH_FALLBACK + + +def test_format_fetch_context_normalizes_non_str_summary() -> None: + out = format_fetch_context_content( + { + "strategy": "search", + "results": [{"item_type": "model", "name": "orders", "summary": 42}], + } + ) + assert "1. [model] orders — 42" in out + + +def test_format_recall_skips_non_dict_rows() -> None: + out = format_recall_content( + [ + "x", + {"nl": "List orders", "sql": "SELECT 1"}, + ] + ) + # Renumbering: the valid row is `1.` even though a bad row preceded it. + assert '1. "List orders"' in out + assert "SELECT 1" in out + + +def test_format_recall_non_list_falls_back() -> None: + assert format_recall_content("nope") == _RECALL_FALLBACK # type: ignore[arg-type] + + +def test_format_recall_all_invalid_falls_back() -> None: + assert format_recall_content(["x", None, 3]) == _RECALL_FALLBACK + + +def test_format_list_models_skips_non_dict_models() -> None: + out = format_list_models_content( + { + "models": [ + "bad", + None, + { + "name": "customers", + "columns": [{}, {}], + "description": "desc", + }, + ] + } + ) + assert "| customers | 2 | desc |" in out + + +def test_format_list_models_non_list_falls_back() -> None: + assert format_list_models_content({"models": "oops"}) == _MODELS_FALLBACK + + +def test_format_list_models_all_invalid_falls_back() -> None: + assert format_list_models_content({"models": ["x", None]}) == _MODELS_FALLBACK + + +def test_format_list_models_normalizes_non_list_columns() -> None: + out = format_list_models_content( + {"models": [{"name": "customers", "columns": "nope", "description": "d"}]} + ) + assert "| customers | 0 | d |" in out + + +def test_format_list_models_normalizes_non_dict_properties() -> None: + out = format_list_models_content( + { + "models": [ + { + "name": "customers", + "columns": [{}], + "properties": "oops", + "description": "fallback desc", + } + ] + } + ) + assert "| customers | 1 | fallback desc |" in out