refac
This commit is contained in:
@@ -260,9 +260,14 @@ async def get_chat_context_usage(chat: Any, model_id: str | None = None) -> dict
|
||||
|
||||
for idx in range(len(messages) - 1, -1, -1):
|
||||
usage = messages[idx].get('usage') or (messages[idx].get('info') or {}).get('usage')
|
||||
input_tokens = (usage or {}).get('input_tokens') or (usage or {}).get('prompt_tokens')
|
||||
input_tokens = (
|
||||
(usage or {}).get('prompt_tokens')
|
||||
or (usage or {}).get('input_tokens')
|
||||
)
|
||||
if isinstance(usage, dict) and input_tokens:
|
||||
tokens = int(input_tokens or 0) + int(usage.get('output_tokens') or usage.get('completion_tokens') or 0)
|
||||
tokens = int(input_tokens or 0) + int(
|
||||
usage.get('completion_tokens') or usage.get('output_tokens') or 0
|
||||
)
|
||||
tokens += _estimate_messages_tokens(messages[idx + 1 :])
|
||||
return _build_context_usage(tokens, threshold)
|
||||
|
||||
@@ -301,8 +306,10 @@ def _exceeds_token_threshold(messages: list[dict], system_prompt: str, summary:
|
||||
|
||||
for idx in range(len(messages) - 1, -1, -1):
|
||||
usage = messages[idx].get('usage') or (messages[idx].get('info') or {}).get('usage')
|
||||
if isinstance(usage, dict) and usage.get('input_tokens'):
|
||||
total = int(usage.get('input_tokens') or 0) + int(usage.get('output_tokens') or 0)
|
||||
if isinstance(usage, dict) and (usage.get('prompt_tokens') or usage.get('input_tokens')):
|
||||
total = int(usage.get('prompt_tokens') or usage.get('input_tokens') or 0) + int(
|
||||
usage.get('completion_tokens') or usage.get('output_tokens') or 0
|
||||
)
|
||||
return total + _estimate_messages_tokens(messages[idx + 1 :]) > threshold
|
||||
|
||||
estimated = _estimate_tokens(system_prompt) + _estimate_tokens(summary or '') + _estimate_messages_tokens(messages)
|
||||
|
||||
Reference in New Issue
Block a user