mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-27 04:20:09 +00:00
zeekay
ea3116aac9
- Remove all unconditional pytest.skip() calls - Fix file content tests to handle both strict/non-strict modes - Fix missing device_capabilities module in hanzo-network - Fix HanzoNetProvider imports and exports - Fix local_llm.py import errors - Create test helper classes for pagination tests - Fix FastMCP paginator test constructor args - Configure all API keys from ~/.zsh/after.zsh - All tests now pass or skip conditionally based on actual requirements
102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
"""Response truncation utilities for MCP tools.
|
|
|
|
This module provides utilities to ensure MCP tool responses don't exceed token limits.
|
|
"""
|
|
|
|
import tiktoken
|
|
|
|
|
|
def estimate_tokens(text: str, model: str = "gpt-4") -> int:
|
|
"""Estimate the number of tokens in a text string.
|
|
|
|
Args:
|
|
text: The text to estimate tokens for
|
|
model: The model to use for token estimation (default: gpt-4)
|
|
|
|
Returns:
|
|
Estimated number of tokens
|
|
"""
|
|
try:
|
|
# Try to get the encoding for the specific model
|
|
encoding = tiktoken.encoding_for_model(model)
|
|
except KeyError:
|
|
# Fall back to cl100k_base which is used by newer models
|
|
encoding = tiktoken.get_encoding("cl100k_base")
|
|
|
|
return len(encoding.encode(text))
|
|
|
|
|
|
def truncate_response(
|
|
response: str,
|
|
max_tokens: int = 20000,
|
|
truncation_message: str = "\n\n[Response truncated due to length. Please use pagination, filtering, or limit parameters to see more.]",
|
|
) -> str:
|
|
"""Truncate a response to fit within token limits.
|
|
|
|
Args:
|
|
response: The response text to truncate
|
|
max_tokens: Maximum number of tokens allowed (default: 20000)
|
|
truncation_message: Message to append when truncating
|
|
|
|
Returns:
|
|
Truncated response if needed, original response otherwise
|
|
"""
|
|
# Quick check - if response is short, no need to count tokens
|
|
if len(response) < max_tokens * 2: # Rough estimate: 1 token ≈ 2-4 chars
|
|
return response
|
|
|
|
# Estimate tokens
|
|
token_count = estimate_tokens(response)
|
|
|
|
# If within limit, return as-is
|
|
if token_count <= max_tokens:
|
|
return response
|
|
|
|
# Need to truncate
|
|
# Binary search to find the right truncation point
|
|
left, right = 0, len(response)
|
|
truncation_msg_tokens = estimate_tokens(truncation_message)
|
|
target_tokens = max_tokens - truncation_msg_tokens
|
|
|
|
while left < right - 1:
|
|
mid = (left + right) // 2
|
|
mid_tokens = estimate_tokens(response[:mid])
|
|
|
|
if mid_tokens <= target_tokens:
|
|
left = mid
|
|
else:
|
|
right = mid
|
|
|
|
# Find a good break point (newline or space)
|
|
truncate_at = left
|
|
for i in range(min(100, left), -1, -1):
|
|
if response[left - i] in "\n ":
|
|
truncate_at = left - i
|
|
break
|
|
|
|
return response[:truncate_at] + truncation_message
|
|
|
|
|
|
def truncate_lines(
|
|
response: str,
|
|
max_lines: int = 1000,
|
|
truncation_message: str = "\n\n[Response truncated to {max_lines} lines. Please use pagination or filtering to see more.]",
|
|
) -> str:
|
|
"""Truncate a response by number of lines.
|
|
|
|
Args:
|
|
response: The response text to truncate
|
|
max_lines: Maximum number of lines allowed (default: 1000)
|
|
truncation_message: Message template to append when truncating
|
|
|
|
Returns:
|
|
Truncated response if needed, original response otherwise
|
|
"""
|
|
lines = response.split("\n")
|
|
|
|
if len(lines) <= max_lines:
|
|
return response
|
|
|
|
truncated = "\n".join(lines[:max_lines])
|
|
return truncated + truncation_message.format(max_lines=max_lines)
|