|
6 | 6 |
|
7 | 7 | from agent_kernel import Firewall |
8 | 8 | from agent_kernel.firewall.budgets import Budgets |
| 9 | +from agent_kernel.firewall.summarize import summarize |
9 | 10 | from agent_kernel.models import Handle, RawResult |
10 | 11 |
|
11 | 12 |
|
@@ -155,3 +156,133 @@ def test_max_depth_limiting() -> None: |
155 | 156 | budgets = Budgets(max_depth=2) |
156 | 157 | frame = _transform(deep, "summary", budgets=budgets) |
157 | 158 | assert frame.response_mode == "summary" # type: ignore[union-attr] |
| 159 | + |
| 160 | + |
| 161 | +# ── Raw mode budget warning ──────────────────────────────────────────────────── |
| 162 | + |
| 163 | + |
| 164 | +def test_raw_mode_oversized_data_adds_warning() -> None: |
| 165 | + large_data = {"payload": "x" * 10_000} |
| 166 | + budgets = Budgets(max_chars=100) |
| 167 | + frame = _transform(large_data, "raw", principal_roles=["admin"], budgets=budgets) |
| 168 | + assert frame.response_mode == "raw" # type: ignore[union-attr] |
| 169 | + assert any("exceeds budget" in w for w in frame.warnings) # type: ignore[union-attr] |
| 170 | + assert frame.raw_data == large_data # type: ignore[union-attr] |
| 171 | + |
| 172 | + |
| 173 | +# ── Table mode with non-list data ────────────────────────────────────────────── |
| 174 | + |
| 175 | + |
| 176 | +def test_table_mode_single_dict() -> None: |
| 177 | + frame = _transform({"a": 1, "b": 2}, "table") |
| 178 | + assert frame.response_mode == "table" # type: ignore[union-attr] |
| 179 | + assert len(frame.table_preview) == 1 # type: ignore[union-attr] |
| 180 | + assert frame.table_preview[0]["a"] == 1 # type: ignore[union-attr] |
| 181 | + |
| 182 | + |
| 183 | +def test_table_mode_non_dict_rows() -> None: |
| 184 | + frame = _transform([1, 2, 3], "table") |
| 185 | + assert frame.response_mode == "table" # type: ignore[union-attr] |
| 186 | + assert frame.table_preview[0] == {"value": 1} # type: ignore[union-attr] |
| 187 | + |
| 188 | + |
| 189 | +def test_table_mode_scalar_data() -> None: |
| 190 | + frame = _transform(42, "table") |
| 191 | + assert frame.response_mode == "table" # type: ignore[union-attr] |
| 192 | + assert frame.table_preview == [{"value": 42}] # type: ignore[union-attr] |
| 193 | + |
| 194 | + |
| 195 | +# ── _cap_facts via public interface ──────────────────────────────────────────── |
| 196 | + |
| 197 | + |
| 198 | +def test_summary_cap_facts_stops_at_budget() -> None: |
| 199 | + # "Keys: key1, key2" (16 chars) fits in max_chars=20; the next fact (46+ chars) |
| 200 | + # pushes the running total over budget, triggering the break in _cap_facts. |
| 201 | + data = {"key1": "v" * 40, "key2": "v" * 40} |
| 202 | + budgets = Budgets(max_chars=20) |
| 203 | + frame = _transform(data, "summary", budgets=budgets) |
| 204 | + assert frame.response_mode == "summary" # type: ignore[union-attr] |
| 205 | + assert len(frame.facts) == 1 # type: ignore[union-attr] |
| 206 | + assert "Keys" in frame.facts[0] # type: ignore[union-attr] |
| 207 | + |
| 208 | + |
| 209 | +def test_cap_facts_all_fit() -> None: |
| 210 | + # Both short facts fit well within a generous budget — no break triggered. |
| 211 | + data = {"a": 1, "b": 2} |
| 212 | + budgets = Budgets(max_chars=10_000) |
| 213 | + frame = _transform(data, "summary", budgets=budgets) |
| 214 | + assert frame.response_mode == "summary" # type: ignore[union-attr] |
| 215 | + assert len(frame.facts) >= 2 # type: ignore[union-attr] |
| 216 | + |
| 217 | + |
| 218 | +# ── summarize() edge cases ───────────────────────────────────────────────────── |
| 219 | + |
| 220 | + |
| 221 | +def test_summarize_plain_list() -> None: |
| 222 | + facts = summarize([1, 2, 3, "hello"]) |
| 223 | + assert facts[0] == "List of 4 items" |
| 224 | + assert "1" in facts[1] |
| 225 | + |
| 226 | + |
| 227 | +def test_summarize_other_type_int() -> None: |
| 228 | + facts = summarize(42) |
| 229 | + assert facts == ["42"] |
| 230 | + |
| 231 | + |
| 232 | +def test_summarize_other_type_none() -> None: |
| 233 | + facts = summarize(None) |
| 234 | + assert facts == ["None"] |
| 235 | + |
| 236 | + |
| 237 | +def test_summarize_string_truncation() -> None: |
| 238 | + long_str = "a" * 600 |
| 239 | + facts = summarize(long_str) |
| 240 | + assert len(facts) == 1 |
| 241 | + assert "600 chars total" in facts[0] |
| 242 | + assert facts[0].startswith("a" * 500) |
| 243 | + |
| 244 | + |
| 245 | +def test_summarize_list_of_dicts_numeric_max_facts() -> None: |
| 246 | + rows = [{"n1": i, "n2": i * 2, "n3": i * 3} for i in range(5)] |
| 247 | + # max_facts=3: "Total rows" + "Top keys" = 2, then 1 numeric fact hits limit |
| 248 | + facts = summarize(rows, max_facts=3) |
| 249 | + assert len(facts) <= 3 |
| 250 | + |
| 251 | + |
| 252 | +def test_summarize_list_of_dicts_categorical_distribution() -> None: |
| 253 | + rows = [{"status": s} for s in ["open", "closed", "open", "pending", "closed"]] |
| 254 | + facts = summarize(rows) |
| 255 | + assert any("distribution" in f for f in facts) |
| 256 | + |
| 257 | + |
| 258 | +def test_summarize_list_of_dicts_no_string_values_in_field() -> None: |
| 259 | + # List values are not strings and not numeric — categorical loop skips them |
| 260 | + rows = [{"items": [1, 2]}, {"items": [3, 4]}, {"items": [5]}] |
| 261 | + facts = summarize(rows) |
| 262 | + assert any("Total rows" in f for f in facts) |
| 263 | + |
| 264 | + |
| 265 | +def test_summarize_list_of_dicts_categorical_max_facts() -> None: |
| 266 | + rows = [{"status": s, "kind": k} for s, k in [("a", "x"), ("b", "y"), ("a", "z"), ("b", "x")]] |
| 267 | + # max_facts=3: "Total rows" + "Top keys" + 1 categorical fact, then break |
| 268 | + facts = summarize(rows, max_facts=3) |
| 269 | + assert len(facts) <= 3 |
| 270 | + |
| 271 | + |
| 272 | +def test_summarize_dict_list_value() -> None: |
| 273 | + data = {"items": [1, 2, 3], "count": 3} |
| 274 | + facts = summarize(data) |
| 275 | + assert any("list of 3 items" in f for f in facts) |
| 276 | + |
| 277 | + |
| 278 | +def test_summarize_dict_other_value_type() -> None: |
| 279 | + # Tuple is not int/float/str/list/dict — falls through to repr() |
| 280 | + data = {"pair": (1, 2), "count": 1} |
| 281 | + facts = summarize(data) |
| 282 | + assert any("(1, 2)" in f for f in facts) |
| 283 | + |
| 284 | + |
| 285 | +def test_summarize_dict_max_facts() -> None: |
| 286 | + data = {"a": 1, "b": 2, "c": 3} |
| 287 | + facts = summarize(data, max_facts=2) |
| 288 | + assert len(facts) <= 2 |
0 commit comments