File size: 5,185 Bytes
f774338
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
"""Tests for `_compact_old_tool_messages`.

The compaction function runs on every model_node turn and is the only thing
keeping older tool outputs from overflowing the context window. The original
implementation head-truncated to 300 chars, which can amputate the exact line
containing the answer. This file tests the summarize-preferred variant: when
an LLM-backed summarizer is supplied, old long tool outputs are replaced by
the summarizer's result; the head-truncation path remains as a fallback for
cases where the summarizer is unavailable or fails.
"""

from __future__ import annotations

from langchain_core.messages import AIMessage, HumanMessage, ToolMessage

from lilith_agent.app import (
    _COMPACT_KEEP_RECENT,
    _COMPACT_MAX_CHARS,
    _COMPACT_SUMMARY_PREFIX,
    _compact_old_tool_messages,
)


def _long_tool_msg(name: str, content: str) -> ToolMessage:
    return ToolMessage(tool_call_id=f"tc-{name}", name=name, content=content)


def test_short_tool_messages_pass_through_unchanged():
    msgs = [
        HumanMessage("q"),
        _long_tool_msg("web_search", "short result"),
        AIMessage("ok"),
    ]
    out = _compact_old_tool_messages(msgs)
    assert out[1].content == "short result"


def test_recent_tool_messages_kept_verbatim_even_when_long():
    long = "X" * (_COMPACT_MAX_CHARS * 10)
    msgs = [_long_tool_msg("web_search", long) for _ in range(_COMPACT_KEEP_RECENT)]
    out = _compact_old_tool_messages(msgs)
    for m in out:
        assert m.content == long


def test_old_long_message_is_summarized_when_summarizer_provided():
    long = "X" * (_COMPACT_MAX_CHARS * 5)
    msgs = [
        _long_tool_msg("web_search", long),  # old
        *[_long_tool_msg("web_search", "short") for _ in range(_COMPACT_KEEP_RECENT)],
    ]

    def fake_summarizer(tool_name: str, content: str) -> str:
        return "SUMMARY_OF_FACTS_42"

    out = _compact_old_tool_messages(msgs, summarize_fn=fake_summarizer)
    assert "SUMMARY_OF_FACTS_42" in out[0].content
    assert out[0].content.startswith(_COMPACT_SUMMARY_PREFIX)
    # The old raw payload is gone — summarization replaced it.
    assert "X" * 1000 not in out[0].content


def test_summarizer_receives_tool_name_and_full_content():
    long = "alpha " * 500
    msgs = [
        _long_tool_msg("arxiv_search", long),
        *[_long_tool_msg("web_search", "short") for _ in range(_COMPACT_KEEP_RECENT)],
    ]
    recorded: dict = {}

    def fake_summarizer(tool_name: str, content: str) -> str:
        recorded["name"] = tool_name
        recorded["len"] = len(content)
        return "ok"

    _compact_old_tool_messages(msgs, summarize_fn=fake_summarizer)
    assert recorded["name"] == "arxiv_search"
    assert recorded["len"] == len(long)


def test_summarizer_failure_falls_back_to_head_truncation():
    long = "Y" * (_COMPACT_MAX_CHARS * 5)
    msgs = [
        _long_tool_msg("web_search", long),
        *[_long_tool_msg("web_search", "short") for _ in range(_COMPACT_KEEP_RECENT)],
    ]

    def broken_summarizer(tool_name: str, content: str) -> str:
        raise RuntimeError("llm offline")

    out = _compact_old_tool_messages(msgs, summarize_fn=broken_summarizer)
    content = out[0].content
    # Fallback marker from the original truncation path
    assert "COMPACTED" in content
    # First `max_chars` preserved verbatim
    assert content.startswith("Y" * _COMPACT_MAX_CHARS)


def test_summarizer_returning_empty_falls_back_to_truncation():
    long = "Z" * (_COMPACT_MAX_CHARS * 5)
    msgs = [
        _long_tool_msg("web_search", long),
        *[_long_tool_msg("web_search", "short") for _ in range(_COMPACT_KEEP_RECENT)],
    ]

    def empty_summarizer(tool_name: str, content: str) -> str:
        return ""

    out = _compact_old_tool_messages(msgs, summarize_fn=empty_summarizer)
    assert "COMPACTED" in out[0].content
    assert out[0].content.startswith("Z" * _COMPACT_MAX_CHARS)


def test_no_summarizer_uses_truncation_fallback():
    """Backwards-compat: the original (no summarize_fn) path must still truncate."""
    long = "W" * (_COMPACT_MAX_CHARS * 5)
    msgs = [
        _long_tool_msg("web_search", long),
        *[_long_tool_msg("web_search", "short") for _ in range(_COMPACT_KEEP_RECENT)],
    ]
    out = _compact_old_tool_messages(msgs)
    assert "COMPACTED" in out[0].content
    assert out[0].content.startswith("W" * _COMPACT_MAX_CHARS)


def test_already_summarized_message_is_not_resummarized():
    """If a prior pass already produced a `[COMPACTED SUMMARY] …` content,
    a second pass must skip it — otherwise we waste a cheap-model call every
    turn on the same already-shrunk payload.
    """
    prior_summary = _COMPACT_SUMMARY_PREFIX + "already shrunk facts"
    msgs = [
        _long_tool_msg("web_search", prior_summary),
        *[_long_tool_msg("web_search", "short") for _ in range(_COMPACT_KEEP_RECENT)],
    ]
    calls = {"n": 0}

    def fake_summarizer(tool_name: str, content: str) -> str:
        calls["n"] += 1
        return "should not run"

    out = _compact_old_tool_messages(msgs, summarize_fn=fake_summarizer)
    assert calls["n"] == 0
    assert out[0].content == prior_summary