mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 04:17:56 +00:00
fix(llm,compression): bound cross-turn Responses chaining and make compression stick
In store mode every user turn chained onto the previous response, so the provider's stored transcript grew without bound (measured: 889k prompt tokens for a 37k-token saved history) while every local guard, the compression pipeline included, measured the saved history. Each chained tool round also re-sent the system message, which the server appends rather than dedupes, and a saved compression point was applied exactly once, in the turn that made it. Chaining is now bounded. A turn starts from the saved history when the previous turn's reported prompt reached the chain budget (default: the model's context window), when the conversation was compressed after that turn was produced, or when OPENAI_RESPONSES_CHAIN_ACROSS_TURNS is off. Chained rounds omit an unchanged system head (hash carried in the persisted Responses state). truncation="auto" is available behind a setting as a backstop against a chain that outgrows the model's window. Compression: a saved point is applied at every turn start; the threshold counts the summary plus the queries after the point instead of the raw history; re-compression summarises only the tail on top of the last point; the mid-execution path marks itself persisted and resets the provider chain so the rebuilt messages are the context; an empty summary is rejected; the visible "[Context Compression Summary]" rows are no longer replayed as history; appending the same point twice is a no-op. Cache hints: a per-user prompt_cache_key and an optional prompt_cache_retention on Responses API calls. Measured on Azure with the same client shape as production (stateless OpenAI client, server-side tools, PDF part): tokens billed on the sixth turn fell from 58k to 35k, tool rounds add tens of tokens instead of ~2.8k, the turn after a compression reused the saved summary in under two seconds instead of re-summarising, and the round after a mid-execution compression started from the compressed context instead of the full stored transcript.
This commit is contained in:
1 parent
4b56e56649
commit
04d358ab69
19 files changed
+1099
-38
No files matched your search
@@ -0,0 +1,246 @@
|
||||
"""Compression must stick: a saved summary is applied on every later turn,
|
||||
re-compression only summarises the tail, an empty summary is a failure, and
|
||||
the visible ``[Context Compression Summary]`` rows are never replayed.
|
||||
|
||||
Background (prod + OSS reproduction, 2026-09-01/03): the turn-start path
|
||||
returned the full raw history whenever the conversation was under the
|
||||
threshold, so a summary was used exactly once; over the threshold it
|
||||
re-summarised everything from query 0 on every turn (14-24 s each); one
|
||||
conversation was "compressed" to a 0-token summary and carried on with
|
||||
nothing.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from application.api.answer.services.compression import CompressionService
|
||||
from application.api.answer.services.compression.orchestrator import (
|
||||
CompressionOrchestrator,
|
||||
)
|
||||
from application.api.answer.services.compression.threshold_checker import (
|
||||
CompressionThresholdChecker,
|
||||
)
|
||||
from application.api.answer.services.compression.token_counter import TokenCounter
|
||||
from application.api.answer.services.compression.types import CompressionResult
|
||||
from application.api.answer.services.conversation_service import (
|
||||
COMPRESSION_SUMMARY_PROMPT,
|
||||
)
|
||||
|
||||
EPOCH = "2026-09-03T09:00:00+00:00"
|
||||
POINT = {
|
||||
"timestamp": EPOCH,
|
||||
"query_index": 1,
|
||||
"compressed_summary": "S",
|
||||
"original_token_count": 900,
|
||||
"compressed_token_count": 5,
|
||||
"compression_ratio": 180.0,
|
||||
"model_used": "m",
|
||||
"compression_prompt_version": "v1.0",
|
||||
}
|
||||
|
||||
|
||||
def _compressed_conversation():
|
||||
big = "word " * 600
|
||||
return {
|
||||
"queries": [
|
||||
{"prompt": "q0", "response": big},
|
||||
{"prompt": "q1", "response": big},
|
||||
{"prompt": "q2", "response": "r2"},
|
||||
{"prompt": "q3", "response": "r3"},
|
||||
],
|
||||
"compression_metadata": {
|
||||
"is_compressed": True,
|
||||
"last_compression_at": EPOCH,
|
||||
"compression_points": [POINT],
|
||||
},
|
||||
"agent_id": "agent-1",
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def conversation_service():
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def threshold_checker():
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def orchestrator(conversation_service, threshold_checker):
|
||||
return CompressionOrchestrator(
|
||||
conversation_service=conversation_service, threshold_checker=threshold_checker
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestTurnStartReuse:
|
||||
def test_under_threshold_returns_existing_summary_and_recent(
|
||||
self, orchestrator, conversation_service, threshold_checker
|
||||
):
|
||||
conversation_service.get_conversation.return_value = _compressed_conversation()
|
||||
threshold_checker.should_compress.return_value = False
|
||||
|
||||
result = orchestrator.compress_if_needed("conv1", "user1", "m", {"sub": "user1"})
|
||||
|
||||
assert result.success is True
|
||||
assert result.compression_performed is False
|
||||
assert result.compressed_summary == "S"
|
||||
assert [q["prompt"] for q in result.recent_queries] == ["q2", "q3"]
|
||||
assert result.last_compression_at == EPOCH
|
||||
|
||||
def test_uncompressed_under_threshold_returns_full_history(
|
||||
self, orchestrator, conversation_service, threshold_checker
|
||||
):
|
||||
conv = _compressed_conversation()
|
||||
conv["compression_metadata"] = {}
|
||||
conversation_service.get_conversation.return_value = conv
|
||||
threshold_checker.should_compress.return_value = False
|
||||
|
||||
result = orchestrator.compress_if_needed("conv1", "user1", "m", {"sub": "user1"})
|
||||
|
||||
assert result.compressed_summary is None
|
||||
assert len(result.recent_queries) == 4
|
||||
assert result.last_compression_at is None
|
||||
|
||||
@patch(
|
||||
"application.api.answer.services.compression.orchestrator.get_provider_from_model_id",
|
||||
return_value="openai",
|
||||
)
|
||||
@patch(
|
||||
"application.api.answer.services.compression.orchestrator.get_api_key_for_provider",
|
||||
return_value="sk",
|
||||
)
|
||||
@patch("application.api.answer.services.compression.orchestrator.LLMCreator")
|
||||
@patch("application.api.answer.services.compression.orchestrator.CompressionService")
|
||||
@patch("application.api.answer.services.compression.orchestrator.settings")
|
||||
def test_over_threshold_compresses_only_the_tail(
|
||||
self,
|
||||
mock_settings,
|
||||
MockCompressionService,
|
||||
MockLLMCreator,
|
||||
_key,
|
||||
_provider,
|
||||
orchestrator,
|
||||
conversation_service,
|
||||
threshold_checker,
|
||||
):
|
||||
mock_settings.COMPRESSION_MODEL_OVERRIDE = None
|
||||
conversation = _compressed_conversation()
|
||||
conversation_service.get_conversation.return_value = conversation
|
||||
threshold_checker.should_compress.return_value = True
|
||||
MockLLMCreator.create_llm.return_value = MagicMock()
|
||||
metadata = MagicMock()
|
||||
metadata.compression_ratio = 3.0
|
||||
metadata.original_token_count = 30
|
||||
metadata.compressed_token_count = 10
|
||||
metadata.timestamp = "2026-09-03T10:00:00+00:00"
|
||||
svc = MagicMock()
|
||||
svc.compress_and_save.return_value = metadata
|
||||
svc.get_compressed_context.return_value = ("S2", [])
|
||||
MockCompressionService.return_value = svc
|
||||
|
||||
result = orchestrator.compress_if_needed("conv1", "user1", "m", {"sub": "user1"})
|
||||
|
||||
assert result.success and result.compression_performed
|
||||
args, kwargs = svc.compress_and_save.call_args
|
||||
# queries 0-1 are already inside point 1; only 2-3 are new.
|
||||
start_index = kwargs.get("start_index", args[3] if len(args) > 3 else 0)
|
||||
assert start_index == 2
|
||||
assert (kwargs.get("compress_up_to_index") or args[2]) == 3
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestEffectiveTokenCount:
|
||||
def test_counts_summary_plus_recent_only(self):
|
||||
conv = _compressed_conversation()
|
||||
effective = TokenCounter.count_effective_conversation_tokens(conv)
|
||||
raw = TokenCounter.count_conversation_tokens(conv)
|
||||
expected = TokenCounter.count_message_tokens([{"content": "S"}]) + (
|
||||
TokenCounter.count_query_tokens(conv["queries"][2:])
|
||||
)
|
||||
assert effective == expected
|
||||
assert effective < raw
|
||||
|
||||
def test_uncompressed_conversation_counts_everything(self):
|
||||
conv = _compressed_conversation()
|
||||
conv["compression_metadata"] = None
|
||||
assert TokenCounter.count_effective_conversation_tokens(conv) == (
|
||||
TokenCounter.count_conversation_tokens(conv)
|
||||
)
|
||||
|
||||
@patch(
|
||||
"application.api.answer.services.compression.threshold_checker.get_token_limit",
|
||||
return_value=1000,
|
||||
)
|
||||
def test_should_compress_uses_effective_count(self, _limit):
|
||||
checker = CompressionThresholdChecker(threshold_percentage=0.8)
|
||||
conv = _compressed_conversation()
|
||||
# Raw history is ~1.2k tokens (over 800); summary + tail is tiny.
|
||||
assert TokenCounter.count_conversation_tokens(conv) > 800
|
||||
assert checker.should_compress(conv, "m", current_query_tokens=10) is False
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestServiceIncremental:
|
||||
def _service(self, summary_text="<summary>new</summary>"):
|
||||
llm = MagicMock()
|
||||
llm.gen.return_value = summary_text
|
||||
svc = CompressionService(llm=llm, model_id="m")
|
||||
svc.prompt_builder = MagicMock(version="v1.0")
|
||||
svc.prompt_builder.build_prompt.return_value = [{"role": "user", "content": "p"}]
|
||||
return svc
|
||||
|
||||
def test_compress_conversation_tail_only(self):
|
||||
svc = self._service()
|
||||
conv = _compressed_conversation()
|
||||
|
||||
metadata = svc.compress_conversation(conv, compress_up_to_index=3, start_index=2)
|
||||
|
||||
queries, existing = svc.prompt_builder.build_prompt.call_args[0]
|
||||
assert [q["prompt"] for q in queries] == ["q2", "q3"]
|
||||
assert existing == [POINT]
|
||||
assert metadata.query_index == 3
|
||||
assert metadata.compressed_summary == "new"
|
||||
|
||||
def test_nothing_new_since_last_point_is_rejected(self):
|
||||
svc = self._service()
|
||||
with pytest.raises(ValueError):
|
||||
svc.compress_conversation(
|
||||
_compressed_conversation(), compress_up_to_index=1, start_index=2
|
||||
)
|
||||
|
||||
def test_empty_summary_raises(self):
|
||||
svc = self._service("<summary> </summary>")
|
||||
with pytest.raises(ValueError):
|
||||
svc.compress_conversation(_compressed_conversation(), compress_up_to_index=3)
|
||||
|
||||
def test_get_compressed_context_skips_summary_rows(self):
|
||||
svc = CompressionService(llm=None, model_id="m")
|
||||
conv = _compressed_conversation()
|
||||
conv["queries"].insert(2, {"prompt": COMPRESSION_SUMMARY_PROMPT, "response": "S"})
|
||||
summary, recent = svc.get_compressed_context(conv)
|
||||
assert summary == "S"
|
||||
assert [q["prompt"] for q in recent] == ["q2", "q3"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestHistoryHelpers:
|
||||
def test_as_history_skips_summary_rows(self):
|
||||
result = CompressionResult.success_no_compression(
|
||||
[
|
||||
{"prompt": "q", "response": "r"},
|
||||
{"prompt": COMPRESSION_SUMMARY_PROMPT, "response": "S"},
|
||||
]
|
||||
)
|
||||
assert [h["prompt"] for h in result.as_history()] == ["q"]
|
||||
|
||||
def test_success_from_existing(self):
|
||||
result = CompressionResult.success_from_existing(
|
||||
"S", [{"prompt": "q", "response": "r"}], last_compression_at=EPOCH
|
||||
)
|
||||
assert result.success and not result.compression_performed
|
||||
assert result.compressed_summary == "S"
|
||||
assert result.last_compression_at == EPOCH
|
||||
@@ -157,7 +157,7 @@ class TestCompressConversation:
|
||||
{"query_index": 0, "compressed_summary": "Previous summary"}
|
||||
]
|
||||
|
||||
def test_zero_compressed_tokens_ratio(self, mock_llm, sample_conversation):
|
||||
def test_zero_compressed_tokens_is_rejected(self, mock_llm, sample_conversation):
|
||||
mock_builder = MagicMock()
|
||||
mock_builder.build_prompt.return_value = [
|
||||
{"role": "system", "content": "C"},
|
||||
@@ -175,8 +175,10 @@ class TestCompressConversation:
|
||||
MockTC.count_query_tokens.return_value = 1000
|
||||
MockTC.count_message_tokens.return_value = 0
|
||||
|
||||
result = svc.compress_conversation(sample_conversation, 2)
|
||||
assert result.compression_ratio == 0
|
||||
# An empty summary must never become a compression point (it
|
||||
# replaced a 494k-token conversation with nothing in prod).
|
||||
with pytest.raises(ValueError, match="empty summary"):
|
||||
svc.compress_conversation(sample_conversation, 2)
|
||||
|
||||
def test_llm_error_propagates(self, sample_conversation):
|
||||
llm = MagicMock()
|
||||
|
||||
Reference in new issue
Block a user