feat: cap conversation history with token budget to prevent unbounded prompt growth #49
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,56 @@ | ||
| """Conversation history budget management. | ||
|
|
||
| Provides token estimation and history trimming logic that operates on | ||
| Gemini chat session history lists — no Gemini SDK dependency needed. | ||
| """ | ||
|
|
||
| import os | ||
| import logging | ||
|
|
||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| MAX_HISTORY_TOKENS = int(os.getenv("MAX_HISTORY_TOKENS", "16000")) | ||
| MAX_HISTORY_TURNS = int(os.getenv("MAX_HISTORY_TURNS", "50")) | ||
|
|
||
|
|
||
| def estimate_tokens(text: str) -> int: | ||
| """Cheap local token estimate — ~4 chars per token on average.""" | ||
| return max(1, len(text) // 4) | ||
|
|
||
|
|
||
| def trim_history(chat_session) -> bool: | ||
| """Enforce token budget and turn cap on chat history. | ||
|
|
||
| Drops oldest turn-pairs (user + model) until both budgets are satisfied. | ||
| Returns True if any turns were dropped. | ||
| """ | ||
| if not chat_session or not chat_session.history: | ||
| return False | ||
|
|
||
| original_len = len(chat_session.history) | ||
| truncated = False | ||
|
|
||
| while len(chat_session.history) > MAX_HISTORY_TURNS * 2: | ||
| chat_session.history.pop(0) | ||
| chat_session.history.pop(0) | ||
| truncated = True | ||
|
|
||
| while len(chat_session.history) >= 2: | ||
| total = sum( | ||
| estimate_tokens(m.parts[0].text) | ||
| if hasattr(m, "parts") and m.parts and hasattr(m.parts[0], "text") | ||
| else 0 | ||
| for m in chat_session.history | ||
|
Comment on lines
+40
to
+44
There was a problem hiding this comment. 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
# Locate the file and inspect the relevant section with line numbers.
fd -a '^history\.py$' .Repository: Deen-Bridge/dnb-ai Length of output: 186 🏁 Script executed: #!/bin/bash
set -euo pipefail
# Find the target file and relevant tests, then inspect the surrounding code.
python3 - <<'PY'
import os
from pathlib import Path
matches = []
for p in Path('.').rglob('history.py'):
matches.append(str(p))
print("HISTORY_FILES:")
for m in matches:
print(m)
print("\nTEST_FILES:")
for p in Path('.').rglob('*test*'):
if p.is_file() and p.suffix in {'.py', '.ts', '.js'}:
if 'history' in p.name.lower() or 'token' in p.name.lower() or 'gemini' in p.name.lower():
print(p)
PYRepository: Deen-Bridge/dnb-ai Length of output: 228 🏁 Script executed: #!/bin/bash
set -euo pipefail
git ls-files | rg '(^|/)history\.py$|test|gemini|token'Repository: Deen-Bridge/dnb-ai Length of output: 518 🏁 Script executed: #!/bin/bash
set -euo pipefail
# Locate the target file and inspect the relevant section with line numbers.
file=$(git ls-files | rg '(^|/)history\.py$' | head -n 1)
echo "FILE=$file"
if [ -n "$file" ]; then
nl -ba "$file" | sed -n '1,140p'
fi
echo
echo "TESTS:"
git ls-files | rg '(^|/).*(test|spec).*\.(py|js|ts)$|(^|/)tests?/' | rg 'history|token|gemini|chat'Repository: Deen-Bridge/dnb-ai Length of output: 212 🏁 Script executed: #!/bin/bash
set -euo pipefail
file=$(git ls-files | rg '(^|/)history\.py$' | head -n 1)
echo "FILE=$file"
nl -ba "$file" | sed -n '1,160p'Repository: Deen-Bridge/dnb-ai Length of output: 212 🏁 Script executed: #!/bin/bash
set -euo pipefail
echo "hello"Repository: Deen-Bridge/dnb-ai Length of output: 162 🏁 Script executed: #!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for path in [Path('history.py'), Path('tests/test_history_truncation.py')]:
print(f"\n### {path} ###")
text = path.read_text()
lines = text.splitlines()
# Print all lines with numbers if reasonably small, otherwise first 220 lines.
start = 1
end = min(len(lines), 260)
for i in range(start, end + 1):
print(f"{i:4d}: {lines[i-1]}")
PYRepository: Deen-Bridge/dnb-ai Length of output: 9875 🏁 Script executed: #!/bin/bash
set -euo pipefail
# Focus on the specific logic around history truncation and token counting.
python3 - <<'PY'
from pathlib import Path
for file in ['history.py', 'tests/test_history_truncation.py']:
p = Path(file)
print(f"\nFILE: {file}")
for i, line in enumerate(p.read_text().splitlines(), 1):
if 'estimate_tokens' in line or 'parts' in line or 'history' in line or 'truncate' in line or 'token' in line:
print(f"{i:4d}: {line}")
PYRepository: Deen-Bridge/dnb-ai Length of output: 4907 🏁 Script executed: #!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for path in [Path('history.py'), Path('tests/test_history_truncation.py')]:
print(f"\n### {path} ###")
lines = path.read_text().splitlines()
for i, line in enumerate(lines, 1):
if i <= 220:
print(f"{i:4d}: {line}")
PYRepository: Deen-Bridge/dnb-ai Length of output: 9875 🏁 Script executed: #!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for file in ('history.py', 'tests/test_history_truncation.py'):
p = Path(file)
print(f"\n=== {file} ===")
for i, line in enumerate(p.read_text().splitlines(), 1):
if i <= 220:
print(f"{i:4d}: {line}")
PYRepository: Deen-Bridge/dnb-ai Length of output: 9875 🏁 Script executed: #!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for file in ('history.py', 'tests/test_history_truncation.py'):
p = Path(file)
print(f"\n=== {file} ===")
lines = p.read_text().splitlines()
for i, line in enumerate(lines, 1):
if i <= 220:
print(f"{i:4d}: {line}")
PYRepository: Deen-Bridge/dnb-ai Length of output: 9875 Count all text parts in history trimming. Lines 40-44 only budget 🤖 Prompt for AI Agents |
||
| ) | ||
| if total <= MAX_HISTORY_TOKENS: | ||
| break | ||
| chat_session.history.pop(0) | ||
| chat_session.history.pop(0) | ||
| truncated = True | ||
|
|
||
| if truncated: | ||
| turns_dropped = (original_len - len(chat_session.history)) // 2 | ||
| logger.info("Trimmed chat history: dropped %d turn pair(s)", turns_dropped) | ||
|
|
||
| return truncated | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,183 @@ | ||
| """Tests for conversation history truncation. | ||
|
|
||
| All tests run offline — no GEMINI_API_KEY needed. | ||
| """ | ||
|
|
||
| from unittest.mock import MagicMock, PropertyMock | ||
|
|
||
| import pytest | ||
|
There was a problem hiding this comment. 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
printf '%s\n' '--- tracked files of interest ---'
git ls-files | rg '(^|/)(requirements(\.txt)?|pyproject\.toml|setup\.py|setup.cfg|tox.ini|Pipfile|poetry\.lock|tests/test_history_truncation\.py)$'
printf '\n%s\n' '--- tests/test_history_truncation.py ---'
cat -n tests/test_history_truncation.py
printf '\n%s\n' '--- manifest search for pytest ---'
rg -n --hidden --glob '!**/.git/**' '\bpytest\b' requirements.txt requirements*.txt pyproject.toml setup.py setup.cfg tox.ini Pipfile poetry.lock 2>/dev/null || trueRepository: Deen-Bridge/dnb-ai Length of output: 8153 Pin 🤖 Prompt for AI AgentsSource: Path instructions |
||
|
|
||
| from history import ( | ||
| MAX_HISTORY_TOKENS, | ||
| MAX_HISTORY_TURNS, | ||
| estimate_tokens, | ||
| trim_history, | ||
| ) | ||
|
|
||
|
|
||
| # --------------------------------------------------------------------------- | ||
| # Token estimation | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| class TestEstimateTokens: | ||
| def test_empty_string(self): | ||
| assert estimate_tokens("") == 1 | ||
|
|
||
| def test_short_text(self): | ||
| assert estimate_tokens("hello") == 1 | ||
|
|
||
| def test_typical_text(self): | ||
| text = "What is the significance of Salah in Islam?" | ||
| expected = max(1, len(text) // 4) | ||
| assert estimate_tokens(text) == expected | ||
|
|
||
| def test_long_text(self): | ||
| text = "A" * 1000 | ||
| assert estimate_tokens(text) == 250 | ||
|
|
||
|
|
||
| # --------------------------------------------------------------------------- | ||
| # Helpers for building fake history entries | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def _make_msg(role: str, text: str): | ||
| part = MagicMock(spec=[]) | ||
| part.text = text | ||
| msg = MagicMock(spec=["role", "parts"]) | ||
| msg.role = role | ||
| msg.parts = [part] | ||
| return msg | ||
|
|
||
|
|
||
| def _make_turn_pair(user_text: str, model_text: str): | ||
| return _make_msg("user", user_text), _make_msg("model", model_text) | ||
|
|
||
|
|
||
| def _make_chat_session(pairs: list[tuple[str, str]]): | ||
| """Build a fake chat session with a mutable history list.""" | ||
| history = [] | ||
| for user_text, model_text in pairs: | ||
| history.append(_make_msg("user", user_text)) | ||
| history.append(_make_msg("model", model_text)) | ||
| chat = MagicMock(spec=["history"]) | ||
| type(chat).history = PropertyMock(return_value=history) | ||
| return chat | ||
|
|
||
|
|
||
| # --------------------------------------------------------------------------- | ||
| # Truncation behaviour | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| class TestTrimHistory: | ||
| def test_no_history_is_noop(self): | ||
| chat = _make_chat_session([]) | ||
| assert trim_history(chat) is False | ||
|
|
||
| def test_single_turn_under_budget_is_noop(self): | ||
| chat = _make_chat_session([("hello", "hi there")]) | ||
| assert trim_history(chat) is False | ||
|
|
||
| def test_truncates_oldest_pairs_when_over_token_budget(self, monkeypatch): | ||
| monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 10) | ||
| monkeypatch.setattr("history.MAX_HISTORY_TURNS", 50) | ||
| pairs = [ | ||
| ("AAAAA AAAAA AA", "BBBBB BBBBB BB"), # ~5+5 = 10 tokens — will be dropped | ||
| ("CCCCC CCCCC CC", "DDDDD DDDDD DD"), # ~5+5 = 10 tokens — will be dropped | ||
| ("EEEEE EEEEE EE", "FFFFF FFFFF FF"), # ~5+5 = 10 tokens — keeps this | ||
| ] | ||
| chat = _make_chat_session(pairs) | ||
| assert trim_history(chat) is True | ||
| assert len(chat.history) == 2 # one turn pair remains | ||
| assert chat.history[0].parts[0].text == "EEEEE EEEEE EE" | ||
| assert chat.history[1].parts[0].text == "FFFFF FFFFF FF" | ||
|
|
||
| def test_truncates_to_turn_cap_first(self, monkeypatch): | ||
| monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 999999) | ||
| monkeypatch.setattr("history.MAX_HISTORY_TURNS", 2) | ||
| pairs = [ | ||
| ("A", "B"), | ||
| ("C", "D"), | ||
| ("E", "F"), | ||
| ("G", "H"), | ||
| ] | ||
| chat = _make_chat_session(pairs) | ||
| assert trim_history(chat) is True | ||
| # MAX_HISTORY_TURNS=2 → at most 4 entries (2 pairs) | ||
| assert len(chat.history) == 4 | ||
| assert chat.history[0].parts[0].text == "E" | ||
| assert chat.history[1].parts[0].text == "F" | ||
| assert chat.history[2].parts[0].text == "G" | ||
| assert chat.history[3].parts[0].text == "H" | ||
|
|
||
| def test_no_orphaned_turns_after_truncation(self, monkeypatch): | ||
| """History must always have pairs — never a lone user or model turn.""" | ||
| monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 10) | ||
| monkeypatch.setattr("history.MAX_HISTORY_TURNS", 50) | ||
| pairs = [ | ||
| ("AAAAA AAAAA AA", "BBBBB BBBBB BB"), | ||
| ("CCCCC CCCCC CC", "DDDDD DDDDD DD"), | ||
| ("EEEEE EEEEE EE", "FFFFF FFFFF FF"), | ||
| ] | ||
| chat = _make_chat_session(pairs) | ||
| trim_history(chat) | ||
| assert len(chat.history) % 2 == 0 | ||
|
|
||
| def test_returns_true_when_truncation_happens(self, monkeypatch): | ||
| monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 3) | ||
| monkeypatch.setattr("history.MAX_HISTORY_TURNS", 50) | ||
| chat = _make_chat_session([("AAAA AAAA", "BBBB BBBB")]) | ||
| assert trim_history(chat) is True | ||
|
|
||
| def test_preserves_recent_context_after_truncation(self, monkeypatch): | ||
| monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 20) | ||
| monkeypatch.setattr("history.MAX_HISTORY_TURNS", 50) | ||
| pairs = [ | ||
| ("A" * 40 + " very old", "B" * 40 + " very old response"), | ||
| ("C" * 40 + " old", "D" * 40 + " old response"), | ||
| ("Recent question", "Recent answer"), | ||
| ] | ||
| chat = _make_chat_session(pairs) | ||
| trim_history(chat) | ||
| remaining = " ".join( | ||
| m.parts[0].text for m in chat.history | ||
| ) | ||
| assert "Recent" in remaining | ||
| assert "Old question 1" not in remaining | ||
|
Comment on lines
+115
to
+148
There was a problem hiding this comment. 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win Assert actual pair ordering and removed content.
🤖 Prompt for AI Agents |
||
|
|
||
|
|
||
| # --------------------------------------------------------------------------- | ||
| # Integration: end-to-end scenario | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| class TestTrimHistoryIntegration: | ||
| def test_long_conversation_triggers_truncation(self, monkeypatch): | ||
| monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 50) | ||
| monkeypatch.setattr("history.MAX_HISTORY_TURNS", 50) | ||
| pairs = [ | ||
| (f"User message {i} " + "X" * 60, f"Model response {i} " + "Y" * 60) | ||
| for i in range(20) | ||
| ] | ||
| chat = _make_chat_session(pairs) | ||
| truncated = trim_history(chat) | ||
| assert truncated is True | ||
| assert len(chat.history) < len(pairs) * 2 | ||
|
|
||
| def test_budget_configurable_via_env(self, monkeypatch): | ||
| monkeypatch.setattr("history.MAX_HISTORY_TOKENS", 1000) | ||
| monkeypatch.setattr("history.MAX_HISTORY_TURNS", 3) | ||
|
Comment on lines
+169
to
+171
There was a problem hiding this comment. 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win Test environment loading rather than patched globals. This test patches already-imported module constants, so it never validates 🤖 Prompt for AI Agents |
||
| pairs = [ | ||
| ("A", "B"), | ||
| ("C", "D"), | ||
| ("E", "F"), | ||
| ("G", "H"), | ||
| ("I", "J"), | ||
| ] | ||
| chat = _make_chat_session(pairs) | ||
| trim_history(chat) | ||
| # 3 turn pairs max → 6 entries | ||
| assert len(chat.history) <= 6 | ||
| assert len(chat.history) % 2 == 0 | ||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reject invalid history limits at startup.
A negative
MAX_HISTORY_TURNSmakes the loop at Line 34 remain true after history is emptied, so the next pair-pop raisesIndexError. Negative token limits also silently discard all history. Parse these settings as non-negative integers and raise a clear configuration error otherwise.Proposed fix
📝 Committable suggestion
🤖 Prompt for AI Agents