44 lines
1.6 KiB
Python
44 lines
1.6 KiB
Python
# utils.py
|
|
import tiktoken
|
|
|
|
def format_bytes(byte_count):
|
|
"""
|
|
Formats an integer of bytes into a human-readable string in B, KB, or MB.
|
|
|
|
Args:
|
|
byte_count: An integer representing the number of bytes.
|
|
|
|
Returns:
|
|
A string formatted as B, KB, or MB with commas and no decimal places.
|
|
"""
|
|
if not isinstance(byte_count, int):
|
|
raise TypeError("Input must be an integer.")
|
|
|
|
if byte_count < 1024:
|
|
# Format as Bytes if less than 1 KB
|
|
return f"{byte_count:,} B"
|
|
elif byte_count < 1024 * 1024:
|
|
# Format as Kilobytes if less than 1 MB
|
|
kb_value = round(byte_count / 1024)
|
|
return f"{kb_value:,} KB"
|
|
else:
|
|
# Format as Megabytes for 1 MB or more
|
|
mb_value = round(byte_count / (1024 * 1024))
|
|
return f"{mb_value:,} MB"
|
|
|
|
def filter_content_for_summarization(content: str) -> str:
|
|
"""Truncates content to a safe number of tokens for the summarization model."""
|
|
MAX_TOKENS = 16384 # Cap content for summarization at 16k tokens for efficiency
|
|
try:
|
|
encoding = tiktoken.get_encoding("cl100k_base")
|
|
tokens = encoding.encode(content)
|
|
if len(tokens) > MAX_TOKENS:
|
|
truncated_tokens = tokens[:MAX_TOKENS]
|
|
return encoding.decode(truncated_tokens)
|
|
else:
|
|
return content
|
|
except Exception as e:
|
|
# Fallback to simple character truncation if tokenization fails
|
|
print(f"Token-based filtering failed: {e}. Falling back to character-based truncation.")
|
|
return content[:MAX_TOKENS * 4] # Rough approximation
|