참고소스 수정본

참고/knowledge_agent-main/utils.py

@@ -0,0 +1,43 @@
+# 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