fix(compressor): ABC compliance — total_tokens, api_mode, logger consistency

This commit is contained in:
0z1-ghb 2026-05-20 01:26:35 +03:00 committed by Teknium
parent 75643a6154
commit 8b2adead78
5 changed files with 29 additions and 5 deletions

View file

@ -1429,6 +1429,7 @@ def init_agent(
base_url=agent.base_url,
api_key=getattr(agent, "api_key", ""),
provider=agent.provider,
api_mode=agent.api_mode,
)
if not agent.quiet_mode:
_ra().logger.info("Using context engine: %s", _selected_engine.name)

View file

@ -609,6 +609,7 @@ class ContextCompressor(ContextEngine):
"""Update tracked token usage from API response."""
self.last_prompt_tokens = usage.get("prompt_tokens", 0)
self.last_completion_tokens = usage.get("completion_tokens", 0)
self.last_total_tokens = usage.get("total_tokens", self.last_prompt_tokens + self.last_completion_tokens)
def should_compress(self, prompt_tokens: int = None) -> bool:
"""Check if context exceeds the compression threshold.
@ -897,7 +898,7 @@ class ContextCompressor(ContextEngine):
into the warning log.
"""
self._summary_model_fallen_back = True
logging.warning(
logger.warning(
"Summary model '%s' %s (%s). "
"Falling back to main model '%s' for compression.",
self.summary_model, reason, e, self.model,
@ -1086,7 +1087,7 @@ The user has requested that this compaction PRIORITISE preserving all informatio
# No provider configured — long cooldown, unlikely to self-resolve
self._summary_failure_cooldown_until = time.monotonic() + _SUMMARY_FAILURE_COOLDOWN_SECONDS
self._last_summary_error = "no auxiliary LLM provider configured"
logging.warning("Context compression: no provider available for "
logger.warning("Context compression: no provider available for "
"summary. Middle turns will be dropped without summary "
"for %d seconds.",
_SUMMARY_FAILURE_COOLDOWN_SECONDS)
@ -1182,7 +1183,7 @@ The user has requested that this compaction PRIORITISE preserving all informatio
if len(err_text) > 220:
err_text = err_text[:217].rstrip() + "..."
self._last_summary_error = err_text
logging.warning(
logger.warning(
"Failed to generate context summary: %s. "
"Further summary attempts paused for %d seconds.",
e,

View file

@ -200,6 +200,7 @@ class ContextEngine(ABC):
base_url: str = "",
api_key: str = "",
provider: str = "",
api_mode: str = "",
) -> None:
"""Called when the user switches models or on fallback activation.

View file

@ -0,0 +1,22 @@
"""Test that last_total_tokens is correctly set by ContextCompressor."""
from agent.context_compressor import ContextCompressor
def test_update_from_response_sets_total_tokens():
"""ABC contract: last_total_tokens must be set from API response."""
c = ContextCompressor(model="test", quiet_mode=True, config_context_length=200000)
c.update_from_response({"prompt_tokens": 100, "completion_tokens": 30, "total_tokens": 130})
assert c.last_total_tokens == 130
c.update_from_response({"prompt_tokens": 100, "completion_tokens": 30})
assert c.last_total_tokens == 130
def test_session_reset_clears_total_tokens():
"""on_session_reset must zero total_tokens."""
c = ContextCompressor(model="test", quiet_mode=True, config_context_length=200000)
c.update_from_response({"prompt_tokens": 100, "completion_tokens": 30, "total_tokens": 130})
c.on_session_reset()
assert c.last_total_tokens == 0

View file

@ -87,5 +87,4 @@ def test_plugin_engine_update_model_args():
assert kw["context_length"] == 131_072
assert "model" in kw
assert "provider" in kw
# Should NOT pass api_mode — the ABC doesn't accept it
assert "api_mode" not in kw
assert "api_mode" in kw