mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| bb02091aa3 |
@@ -94,6 +94,12 @@ class BasicMemoryConfig(BaseSettings):
|
||||
gt=0,
|
||||
)
|
||||
|
||||
streaming_checksum_threshold_mb: int = Field(
|
||||
default=1,
|
||||
description="File size threshold in MB for using streaming checksum computation. Files larger than this will be processed in chunks to reduce memory usage. Default: 1MB",
|
||||
gt=0,
|
||||
)
|
||||
|
||||
kebab_filenames: bool = Field(
|
||||
default=False,
|
||||
description="Format for generated filenames. False preserves spaces and special chars, True converts them to hyphens for consistency with permalinks",
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Utilities for file operations."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
import re
|
||||
@@ -12,6 +13,27 @@ from loguru import logger
|
||||
from basic_memory.utils import FilePath
|
||||
|
||||
|
||||
def get_streaming_checksum_threshold() -> int:
|
||||
"""Get the streaming checksum threshold from config.
|
||||
|
||||
Returns threshold in bytes. Defaults to 1MB if config is unavailable.
|
||||
"""
|
||||
try:
|
||||
from basic_memory.config import ConfigManager
|
||||
|
||||
config = ConfigManager().config
|
||||
return config.streaming_checksum_threshold_mb * 1024 * 1024
|
||||
except Exception:
|
||||
# Default to 1MB if config unavailable (e.g., during tests)
|
||||
return 1024 * 1024
|
||||
|
||||
|
||||
# Default threshold for streaming vs in-memory checksum computation (1MB)
|
||||
STREAMING_CHECKSUM_THRESHOLD = 1024 * 1024 # 1MB in bytes
|
||||
# Chunk size for streaming reads (64KB - optimal for most file systems)
|
||||
STREAMING_CHUNK_SIZE = 64 * 1024 # 64KB
|
||||
|
||||
|
||||
class FileError(Exception):
|
||||
"""Base exception for file operations."""
|
||||
|
||||
@@ -52,6 +74,42 @@ async def compute_checksum(content: Union[str, bytes]) -> str:
|
||||
raise FileError(f"Failed to compute checksum: {e}")
|
||||
|
||||
|
||||
async def compute_checksum_streaming(file_path: Path, chunk_size: int = STREAMING_CHUNK_SIZE) -> str:
|
||||
"""
|
||||
Compute SHA-256 checksum using streaming for large files.
|
||||
|
||||
This function reads the file in chunks to avoid loading the entire file into memory,
|
||||
which is critical for large binary files (PDFs, images, videos).
|
||||
|
||||
Args:
|
||||
file_path: Path to file to hash
|
||||
chunk_size: Size of chunks to read (default 64KB)
|
||||
|
||||
Returns:
|
||||
SHA-256 hex digest
|
||||
|
||||
Raises:
|
||||
FileError: If checksum computation fails
|
||||
"""
|
||||
try:
|
||||
hasher = hashlib.sha256()
|
||||
|
||||
def read_chunks():
|
||||
"""Read file in chunks synchronously (for thread pool execution)."""
|
||||
with open(file_path, "rb") as f:
|
||||
while chunk := f.read(chunk_size):
|
||||
hasher.update(chunk)
|
||||
|
||||
# Run blocking I/O in thread pool to avoid blocking event loop
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, read_chunks)
|
||||
|
||||
return hasher.hexdigest()
|
||||
except Exception as e: # pragma: no cover
|
||||
logger.error(f"Failed to compute streaming checksum for {file_path}: {e}")
|
||||
raise FileError(f"Failed to compute streaming checksum: {e}")
|
||||
|
||||
|
||||
async def ensure_directory(path: FilePath) -> None:
|
||||
"""
|
||||
Ensure directory exists, creating if necessary.
|
||||
|
||||
@@ -16,7 +16,11 @@ from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from basic_memory import db
|
||||
from basic_memory.config import BasicMemoryConfig, ConfigManager
|
||||
from basic_memory.file_utils import has_frontmatter
|
||||
from basic_memory.file_utils import (
|
||||
has_frontmatter,
|
||||
compute_checksum_streaming,
|
||||
get_streaming_checksum_threshold,
|
||||
)
|
||||
from basic_memory.ignore_utils import load_bmignore_patterns, should_ignore_path
|
||||
from basic_memory.markdown import EntityParser, MarkdownProcessor
|
||||
from basic_memory.models import Entity, Project
|
||||
@@ -154,6 +158,8 @@ class SyncService:
|
||||
"""Compute file checksum in thread pool to avoid blocking the event loop.
|
||||
|
||||
Uses semaphore to limit concurrent file reads and prevent OOM on large projects.
|
||||
For files larger than STREAMING_CHECKSUM_THRESHOLD (1MB), uses streaming
|
||||
checksum computation to avoid loading entire file into memory.
|
||||
"""
|
||||
|
||||
def _sync_compute_checksum(path_str: str) -> str:
|
||||
@@ -175,6 +181,27 @@ class SyncService:
|
||||
return hashlib.sha256(content_bytes).hexdigest()
|
||||
|
||||
async with self._file_semaphore:
|
||||
path_obj = self.file_service.base_path / path
|
||||
|
||||
# Check file size to decide between streaming and in-memory checksum
|
||||
try:
|
||||
file_size = path_obj.stat().st_size
|
||||
threshold = get_streaming_checksum_threshold()
|
||||
|
||||
# For large files (>threshold), use streaming checksum to avoid memory issues
|
||||
if file_size > threshold:
|
||||
logger.debug(
|
||||
f"Using streaming checksum for large file: path={path}, "
|
||||
f"size_mb={file_size / (1024 * 1024):.2f}, "
|
||||
f"threshold_mb={threshold / (1024 * 1024):.2f}"
|
||||
)
|
||||
return await compute_checksum_streaming(path_obj)
|
||||
|
||||
except Exception as e:
|
||||
# If we can't get file size, fall back to original method
|
||||
logger.warning(f"Failed to check file size for {path}: {e}, using non-streaming checksum")
|
||||
|
||||
# For smaller files, use the original in-memory method
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(self._thread_pool, _sync_compute_checksum, path)
|
||||
|
||||
|
||||
@@ -11,6 +11,7 @@ from basic_memory.file_utils import (
|
||||
FileWriteError,
|
||||
ParseError,
|
||||
compute_checksum,
|
||||
compute_checksum_streaming,
|
||||
ensure_directory,
|
||||
has_frontmatter,
|
||||
parse_frontmatter,
|
||||
@@ -19,6 +20,7 @@ from basic_memory.file_utils import (
|
||||
sanitize_for_folder,
|
||||
update_frontmatter,
|
||||
write_file_atomic,
|
||||
STREAMING_CHECKSUM_THRESHOLD,
|
||||
)
|
||||
|
||||
|
||||
@@ -327,3 +329,97 @@ def test_sanitize_for_filename_removes_invalid_characters():
|
||||
)
|
||||
def test_sanitize_for_folder_edge_cases(input_folder, expected):
|
||||
assert sanitize_for_folder(input_folder) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compute_checksum_streaming_small_file(tmp_path: Path):
|
||||
"""Test streaming checksum computation on a small file."""
|
||||
test_file = tmp_path / "small.txt"
|
||||
content = "test content for streaming"
|
||||
test_file.write_text(content, encoding="utf-8")
|
||||
|
||||
# Compute checksum using streaming
|
||||
streaming_checksum = await compute_checksum_streaming(test_file)
|
||||
|
||||
# Compute checksum using regular method for comparison
|
||||
regular_checksum = await compute_checksum(content)
|
||||
|
||||
# Both methods should produce the same checksum
|
||||
assert streaming_checksum == regular_checksum
|
||||
assert len(streaming_checksum) == 64 # SHA-256 produces 64 char hex string
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compute_checksum_streaming_large_file(tmp_path: Path):
|
||||
"""Test streaming checksum computation on a large file (>1MB)."""
|
||||
test_file = tmp_path / "large.bin"
|
||||
|
||||
# Create a file larger than STREAMING_CHECKSUM_THRESHOLD (1MB)
|
||||
# Use 2MB to ensure we're testing the streaming path
|
||||
large_content = b"x" * (2 * 1024 * 1024) # 2MB of 'x' bytes
|
||||
test_file.write_bytes(large_content)
|
||||
|
||||
# Compute checksum using streaming
|
||||
streaming_checksum = await compute_checksum_streaming(test_file)
|
||||
|
||||
# Compute checksum using regular method for comparison
|
||||
regular_checksum = await compute_checksum(large_content)
|
||||
|
||||
# Both methods should produce the same checksum
|
||||
assert streaming_checksum == regular_checksum
|
||||
assert len(streaming_checksum) == 64
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compute_checksum_streaming_binary_file(tmp_path: Path):
|
||||
"""Test streaming checksum on binary file content."""
|
||||
test_file = tmp_path / "binary.bin"
|
||||
|
||||
# Create binary content with various byte values
|
||||
binary_content = bytes(range(256)) * 1000 # ~256KB of binary data
|
||||
test_file.write_bytes(binary_content)
|
||||
|
||||
# Compute checksum using streaming
|
||||
streaming_checksum = await compute_checksum_streaming(test_file)
|
||||
|
||||
# Compute checksum using regular method for comparison
|
||||
regular_checksum = await compute_checksum(binary_content)
|
||||
|
||||
# Both methods should produce the same checksum
|
||||
assert streaming_checksum == regular_checksum
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compute_checksum_streaming_empty_file(tmp_path: Path):
|
||||
"""Test streaming checksum on an empty file."""
|
||||
test_file = tmp_path / "empty.txt"
|
||||
test_file.write_text("", encoding="utf-8")
|
||||
|
||||
# Compute checksum using streaming
|
||||
streaming_checksum = await compute_checksum_streaming(test_file)
|
||||
|
||||
# Compute checksum using regular method for comparison
|
||||
regular_checksum = await compute_checksum("")
|
||||
|
||||
# Both methods should produce the same checksum
|
||||
assert streaming_checksum == regular_checksum
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compute_checksum_streaming_custom_chunk_size(tmp_path: Path):
|
||||
"""Test streaming checksum with custom chunk size."""
|
||||
test_file = tmp_path / "custom.txt"
|
||||
content = "test content " * 10000 # ~130KB
|
||||
test_file.write_text(content, encoding="utf-8")
|
||||
|
||||
# Test with various chunk sizes
|
||||
chunk_sizes = [1024, 4096, 65536] # 1KB, 4KB, 64KB
|
||||
|
||||
checksums = []
|
||||
for chunk_size in chunk_sizes:
|
||||
checksum = await compute_checksum_streaming(test_file, chunk_size=chunk_size)
|
||||
checksums.append(checksum)
|
||||
|
||||
# All chunk sizes should produce the same checksum
|
||||
assert len(set(checksums)) == 1
|
||||
assert checksums[0] == await compute_checksum(content)
|
||||
|
||||
Reference in New Issue
Block a user