Soulter
241f1c26d3
* feat: context compressor Co-authored-by: kawayiYokami <289104862@qq.com> * Add comprehensive tests for ContextManager and ContextTruncator - Implemented a full test suite for ContextManager covering initialization, message processing, token-based compression, and error handling. - Added tests for ContextTruncator focusing on message fixing, truncation by turns, dropping oldest turns, and halving. - Ensured that both test suites validate edge cases and maintain expected behavior with various message types, including system and tool messages. * feat: add MockProvider for LLM compression tests * chore: remove lock * ruff fix * fix * perf * feat: enhance context compression with token tracking and logging * feat: update logging for context compression trigger * feat: implement context compression logic with dynamic threshold and token tracking * fix: reorder import statements for consistency * feat: add token_usage tracking to conversations and update related processing logic --------- Co-authored-by: kawayiYokami <289104862@qq.com>
65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
import json
|
|
from typing import Protocol, runtime_checkable
|
|
|
|
from ..message import Message, TextPart
|
|
|
|
|
|
@runtime_checkable
|
|
class TokenCounter(Protocol):
|
|
"""
|
|
Protocol for token counters.
|
|
Provides an interface for counting tokens in message lists.
|
|
"""
|
|
|
|
def count_tokens(
|
|
self, messages: list[Message], trusted_token_usage: int = 0
|
|
) -> int:
|
|
"""Count the total tokens in the message list.
|
|
|
|
Args:
|
|
messages: The message list.
|
|
trusted_token_usage: The total token usage that LLM API returned.
|
|
For some cases, this value is more accurate.
|
|
But some API does not return it, so the value defaults to 0.
|
|
|
|
Returns:
|
|
The total token count.
|
|
"""
|
|
...
|
|
|
|
|
|
class EstimateTokenCounter:
|
|
"""Estimate token counter implementation.
|
|
Provides a simple estimation of token count based on character types.
|
|
"""
|
|
|
|
def count_tokens(
|
|
self, messages: list[Message], trusted_token_usage: int = 0
|
|
) -> int:
|
|
if trusted_token_usage > 0:
|
|
return trusted_token_usage
|
|
|
|
total = 0
|
|
for msg in messages:
|
|
content = msg.content
|
|
if isinstance(content, str):
|
|
total += self._estimate_tokens(content)
|
|
elif isinstance(content, list):
|
|
# 处理多模态内容
|
|
for part in content:
|
|
if isinstance(part, TextPart):
|
|
total += self._estimate_tokens(part.text)
|
|
|
|
# 处理 Tool Calls
|
|
if msg.tool_calls:
|
|
for tc in msg.tool_calls:
|
|
tc_str = json.dumps(tc if isinstance(tc, dict) else tc.model_dump())
|
|
total += self._estimate_tokens(tc_str)
|
|
|
|
return total
|
|
|
|
def _estimate_tokens(self, text: str) -> int:
|
|
chinese_count = len([c for c in text if "\u4e00" <= c <= "\u9fff"])
|
|
other_count = len(text) - chinese_count
|
|
return int(chinese_count * 0.6 + other_count * 0.3)
|