refac: kb sync
This commit is contained in:
Timothy Jaeryang Baek
@@ -1,4 +1,5 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
@@ -290,6 +291,10 @@ async def upload_file_handler(
|
||||
},
|
||||
)
|
||||
|
||||
# SHA-256 of raw uploaded bytes for incremental sync diffing.
|
||||
# If the client pre-computed and sent file_hash, use that.
|
||||
file_hash = file_metadata.get('file_hash') or hashlib.sha256(contents).hexdigest()
|
||||
|
||||
file_item = await Files.insert_new_file(
|
||||
user.id,
|
||||
FileForm(
|
||||
@@ -304,6 +309,7 @@ async def upload_file_handler(
|
||||
'name': name,
|
||||
'content_type': (file.content_type if isinstance(file.content_type, str) else None),
|
||||
'size': len(contents),
|
||||
'file_hash': file_hash,
|
||||
'data': file_metadata,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -984,6 +984,184 @@ async def reset_knowledge_by_id(
|
||||
return knowledge
|
||||
|
||||
|
||||
############################
|
||||
# SyncKnowledgeDiff
|
||||
############################
|
||||
|
||||
|
||||
class FileManifestEntry(BaseModel):
|
||||
filename: str # basename: "readme.md"
|
||||
path: str # relative dir: "docs/api" or "" for root
|
||||
checksum: str # SHA-256 of raw bytes
|
||||
size: int
|
||||
|
||||
|
||||
class SyncDiffForm(BaseModel):
|
||||
manifest: list[FileManifestEntry]
|
||||
|
||||
|
||||
class SyncDiffResponse(BaseModel):
|
||||
added: list[dict] # [{filename, path}] — new files
|
||||
modified: list[dict] # [{filename, path, stale_file_id}] — changed files
|
||||
deleted: list[dict] # [{file_id, filename}] — files to remove
|
||||
mkdir: list[str] # directory paths to create
|
||||
rmdir: list[str] # directory IDs to remove
|
||||
unmodified_count: int
|
||||
|
||||
|
||||
@router.post('/{id}/sync/diff', response_model=SyncDiffResponse)
|
||||
async def sync_knowledge_diff(
|
||||
id: str,
|
||||
form_data: SyncDiffForm,
|
||||
user=Depends(get_verified_user),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
"""
|
||||
Compare a local file manifest against the knowledge base to determine
|
||||
which files need uploading, removing, and which directories to create/remove.
|
||||
"""
|
||||
await _verify_knowledge_write_access(id, user, db)
|
||||
|
||||
# ── Index existing state ──
|
||||
knowledge_files = await Knowledges.get_files_with_directory_ids(id, db=db)
|
||||
existing_directories = await Knowledges.get_all_directories(id, db=db)
|
||||
|
||||
# Build directory path lookups
|
||||
directory_path_by_id: dict[str, str] = {}
|
||||
directory_id_by_path: dict[str, str] = {}
|
||||
for directory in existing_directories:
|
||||
segments = [directory.name]
|
||||
parent_id = directory.parent_id
|
||||
while parent_id:
|
||||
parent = next((d for d in existing_directories if d.id == parent_id), None)
|
||||
if not parent:
|
||||
break
|
||||
segments.insert(0, parent.name)
|
||||
parent_id = parent.parent_id
|
||||
full_path = '/'.join(segments)
|
||||
directory_path_by_id[directory.id] = full_path
|
||||
directory_id_by_path[full_path] = directory.id
|
||||
|
||||
# Index existing files by (path, filename) → {file_id, checksum}
|
||||
indexed_files: dict[tuple[str, str], dict] = {}
|
||||
for file_model, directory_id in knowledge_files:
|
||||
file_path = directory_path_by_id.get(directory_id, '') if directory_id else ''
|
||||
stored_checksum = (file_model.meta or {}).get('file_hash')
|
||||
indexed_files[(file_path, file_model.filename)] = {
|
||||
'file_id': file_model.id,
|
||||
'checksum': stored_checksum,
|
||||
}
|
||||
|
||||
# ── Diff files ──
|
||||
added: list[dict] = []
|
||||
modified: list[dict] = []
|
||||
deleted: list[dict] = []
|
||||
unmodified_count = 0
|
||||
manifest_keys: set[tuple[str, str]] = set()
|
||||
|
||||
for entry in form_data.manifest:
|
||||
key = (entry.path, entry.filename)
|
||||
manifest_keys.add(key)
|
||||
|
||||
if key not in indexed_files:
|
||||
added.append({'filename': entry.filename, 'path': entry.path})
|
||||
elif indexed_files[key]['checksum'] != entry.checksum:
|
||||
modified.append({
|
||||
'filename': entry.filename,
|
||||
'path': entry.path,
|
||||
'stale_file_id': indexed_files[key]['file_id'],
|
||||
})
|
||||
else:
|
||||
unmodified_count += 1
|
||||
|
||||
for key, file_info in indexed_files.items():
|
||||
if key not in manifest_keys:
|
||||
deleted.append({'file_id': file_info['file_id'], 'filename': key[1]})
|
||||
|
||||
# ── Diff directories ──
|
||||
required_directory_paths: set[str] = set()
|
||||
for entry in form_data.manifest:
|
||||
if entry.path:
|
||||
segments = entry.path.split('/')
|
||||
for depth in range(len(segments)):
|
||||
required_directory_paths.add('/'.join(segments[:depth + 1]))
|
||||
|
||||
mkdir = sorted(
|
||||
[p for p in required_directory_paths if p not in directory_id_by_path],
|
||||
key=lambda p: p.count('/')
|
||||
)
|
||||
|
||||
orphaned_directory_paths = set(directory_id_by_path) - required_directory_paths
|
||||
rmdir = [directory_id_by_path[p] for p in orphaned_directory_paths]
|
||||
|
||||
return SyncDiffResponse(
|
||||
added=added,
|
||||
modified=modified,
|
||||
deleted=deleted,
|
||||
mkdir=mkdir,
|
||||
rmdir=rmdir,
|
||||
unmodified_count=unmodified_count,
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# SyncKnowledgeCleanup
|
||||
############################
|
||||
|
||||
|
||||
class SyncCleanupForm(BaseModel):
|
||||
file_ids: list[str] # file IDs to delete
|
||||
dir_ids: list[str] = [] # directory IDs to rmdir
|
||||
|
||||
|
||||
@router.post('/{id}/sync/cleanup')
|
||||
async def sync_knowledge_cleanup(
|
||||
id: str,
|
||||
form_data: SyncCleanupForm,
|
||||
user=Depends(get_verified_user),
|
||||
db: AsyncSession = Depends(get_async_session),
|
||||
):
|
||||
"""
|
||||
Remove stale files and orphaned directories from a knowledge base
|
||||
after an incremental sync.
|
||||
"""
|
||||
await _verify_knowledge_write_access(id, user, db)
|
||||
|
||||
# ── Remove deleted files ──
|
||||
for file_id in form_data.file_ids:
|
||||
file = await Files.get_file_by_id(file_id, db=db)
|
||||
if not file:
|
||||
continue
|
||||
|
||||
await Knowledges.remove_file_from_knowledge_by_id(id, file_id, db=db)
|
||||
|
||||
try:
|
||||
await ASYNC_VECTOR_DB_CLIENT.delete(
|
||||
collection_name=id, filter={'file_id': file_id}
|
||||
)
|
||||
await ASYNC_VECTOR_DB_CLIENT.delete(
|
||||
collection_name=id, filter={'hash': file.hash}
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
collection_name = f'file-{file_id}'
|
||||
if await ASYNC_VECTOR_DB_CLIENT.has_collection(collection_name):
|
||||
await ASYNC_VECTOR_DB_CLIENT.delete_collection(collection_name)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if file.user_id == user.id or user.role == 'admin':
|
||||
await Files.delete_file_by_id(file_id, db=db)
|
||||
|
||||
# ── Remove orphaned directories (children before parents) ──
|
||||
for dir_id in reversed(form_data.dir_ids):
|
||||
await Knowledges.delete_directory(dir_id, move_files_to_parent=False, db=db)
|
||||
|
||||
return {'status': True}
|
||||
|
||||
|
||||
############################
|
||||
# AddFilesToKnowledge
|
||||
############################
|
||||
|
||||
Reference in New Issue
Block a user