From 74f95a9b0d6b89b241701422170d54e21e701baa Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Thu, 14 May 2026 02:08:35 +0900 Subject: [PATCH] refac --- backend/open_webui/tools/knowledge_fs.py | 141 ++++++++++++++++------- 1 file changed, 99 insertions(+), 42 deletions(-) diff --git a/backend/open_webui/tools/knowledge_fs.py b/backend/open_webui/tools/knowledge_fs.py index 871ee7d082..f055f563ba 100644 --- a/backend/open_webui/tools/knowledge_fs.py +++ b/backend/open_webui/tools/knowledge_fs.py @@ -287,17 +287,22 @@ async def _get_accessible_kb_ids(user: dict, model_knowledge: list[dict] | None, async def _get_accessible_files(user: dict, model_knowledge: list[dict] | None, knowledge_id: str | None = None) -> list[dict]: - """Get all files the user can access, with KB metadata and directory paths.""" + """Get all files the user can access, with KB metadata and directory_id (no path computation).""" from open_webui.models.files import Files + from open_webui.models.knowledge import Knowledges kb_ids = await _get_accessible_kb_ids(user, model_knowledge, knowledge_id) files = [] for kb_id, kb_name in kb_ids: - tree = await _build_directory_tree(kb_id) - for f in tree['files']: + kb_files = await Knowledges.get_files_with_directory_ids(kb_id) + for file_model, dir_id in kb_files: files.append({ - **f, + 'id': file_model.id, 'filename': file_model.filename, + 'directory_id': dir_id, + 'size': file_model.meta.get('size') if file_model.meta else None, + 'type': file_model.meta.get('content_type') if file_model.meta else None, + 'updated_at': file_model.updated_at, 'knowledge_id': kb_id, 'knowledge_name': kb_name, }) @@ -313,7 +318,7 @@ async def _get_accessible_files(user: dict, model_knowledge: list[dict] | None, if f: files.append({ 'id': f.id, 'filename': f.filename, - 'path': f.filename, 'directory_id': None, + 'directory_id': None, 'size': f.meta.get('size') if f.meta else None, 'type': f.meta.get('content_type') if f.meta else None, 'updated_at': f.updated_at, @@ -323,11 +328,44 @@ async def _get_accessible_files(user: dict, model_knowledge: list[dict] | None, return files +async def _resolve_dir_path(path: str, knowledge_id: str) -> str | None: + """Walk a directory path one level at a time. Returns dir_id or None.""" + from open_webui.models.knowledge import Knowledges + + parts = path.strip('/').split('/') + current_parent = None + + for part in parts: + dirs = await Knowledges.get_directories(knowledge_id, parent_id=current_parent) + match = next((d for d in dirs if d.name == part), None) + if not match: + return None + current_parent = match.id + + return current_parent + + +async def _get_descendant_dir_ids(dir_id: str, knowledge_id: str) -> set[str]: + """Collect all descendant directory IDs recursively.""" + from open_webui.models.knowledge import Knowledges + + result = {dir_id} + queue = [dir_id] + while queue: + parent = queue.pop() + children = await Knowledges.get_directories(knowledge_id, parent_id=parent) + for child in children: + if child.id not in result: + result.add(child.id) + queue.append(child.id) + return result + + async def _resolve_file(ref: str, user: dict, model_knowledge: list[dict] | None) -> dict | None: """Resolve a file reference (ID, path, or filename) to a file info dict with content.""" from open_webui.models.files import Files - # Get the set of files this user can access + # Get accessible file IDs (lightweight — no path computation) accessible = await _get_accessible_files(user, model_knowledge) accessible_ids = {fi['id'] for fi in accessible} @@ -335,24 +373,30 @@ async def _resolve_file(ref: str, user: dict, model_knowledge: list[dict] | None f = await Files.get_file_by_id(ref) if f and f.data: if f.id not in accessible_ids: - return None # User cannot access this file + return None return {'id': f.id, 'filename': f.filename, 'content': f.data.get('content', ''), 'meta': f.meta, 'updated_at': f.updated_at, 'created_at': f.created_at} - # Try path match (e.g. "docs/api/auth.md") + # Try path match (e.g. "docs/api/auth.md") — lazy dir walk ref_clean = ref.strip('/') if '/' in ref_clean: - path_matches = [fi for fi in accessible if fi.get('path') == ref_clean] - if len(path_matches) == 1: - f = await Files.get_file_by_id(path_matches[0]['id']) - if f and f.data: - return {'id': f.id, 'filename': f.filename, 'content': f.data.get('content', ''), - 'meta': f.meta, 'updated_at': f.updated_at, 'created_at': f.created_at, - 'knowledge_id': path_matches[0].get('knowledge_id'), - 'knowledge_name': path_matches[0].get('knowledge_name')} - elif len(path_matches) > 1: - return {'error': f'Ambiguous path "{ref}". Matches:\n' + - '\n'.join(f" {m['id']} {m['path']} ({m.get('knowledge_name', 'direct')})" for m in path_matches)} + dir_path, filename = ref_clean.rsplit('/', 1) + # Try resolving in each accessible KB + kb_ids = {fi['knowledge_id'] for fi in accessible if fi.get('knowledge_id')} + for kb_id in kb_ids: + dir_id = await _resolve_dir_path(dir_path, kb_id) + if dir_id is None: + continue + # Find file with that name in that directory + matches = [fi for fi in accessible + if fi['filename'] == filename and fi['directory_id'] == dir_id] + if len(matches) == 1: + f = await Files.get_file_by_id(matches[0]['id']) + if f and f.data: + return {'id': f.id, 'filename': f.filename, 'content': f.data.get('content', ''), + 'meta': f.meta, 'updated_at': f.updated_at, 'created_at': f.created_at, + 'knowledge_id': matches[0].get('knowledge_id'), + 'knowledge_name': matches[0].get('knowledge_name')} # Try filename match within accessible files matches = [fi for fi in accessible if fi['filename'] == ref] @@ -366,7 +410,7 @@ async def _resolve_file(ref: str, user: dict, model_knowledge: list[dict] | None 'knowledge_name': matches[0].get('knowledge_name')} elif len(matches) > 1: return {'error': f'Ambiguous filename "{ref}". Use full path to disambiguate:\n' + - '\n'.join(f" {m['id']} {m.get('path', m['filename'])} ({m.get('knowledge_name', 'direct')})" for m in matches)} + '\n'.join(f" {m['id']} {m['filename']} ({m.get('knowledge_name', 'direct')})" for m in matches)} return None @@ -388,10 +432,11 @@ async def _get_file_content(file_id: str) -> str | None: async def _kb_ls(args: list[str], flags: set[str], user: dict, model_knowledge: list[dict] | None) -> str: """List files and directories. Supports: ls, ls , ls -a (flat).""" + from open_webui.models.knowledge import Knowledges + flat_mode = 'a' in flags path_arg = args[0] if args else None - # Determine which KBs are accessible kb_ids = await _get_accessible_kb_ids(user, model_knowledge, knowledge_id=None) # If path_arg looks like a KB ID, scope to that KB @@ -413,36 +458,34 @@ async def _kb_ls(args: list[str], flags: set[str], user: dict, lines = [] for kb_id, kb_name in kb_ids: - tree = await _build_directory_tree(kb_id) lines.append(f'Knowledge Base: {kb_name} ({kb_id})') if flat_mode: - # Flat: show all files with full paths + # Flat mode: build full tree (legitimate use) + tree = await _build_directory_tree(kb_id) for f in tree['files']: lines.append(f' {f["id"]} {f["path"]} {_fmt_size(f)} {_fmt_date(f)}') lines.append('') continue - # Resolve target directory + # Resolve target directory (lazy walk) target_dir_id = None if dir_path: - target_dir_id = _resolve_path(dir_path, tree) + target_dir_id = await _resolve_dir_path(dir_path, kb_id) if target_dir_id is None: lines.append(f' Directory not found: {dir_path}') lines.append('') continue lines.append(f' Path: {dir_path}/') - # Show subdirectories - subdirs = _get_subdirs(tree, target_dir_id) + # Show subdirectories (targeted query — only this level) + subdirs = await Knowledges.get_directories(kb_id, parent_id=target_dir_id) for d in subdirs: - # Count files recursively under this dir - child_files = _get_files_under_dir(tree, d['id']) - count_str = f'{len(child_files)} file{"s" if len(child_files) != 1 else ""}' - lines.append(f' 📁 {d["name"]}/ ({count_str})') + lines.append(f' 📁 {d.name}/') - # Show files at this level - dir_files = _get_files_in_dir(tree, target_dir_id) + # Show files at this level (filter from accessible files) + accessible = await _get_accessible_files(user, model_knowledge, knowledge_id=kb_id) + dir_files = [f for f in accessible if f['directory_id'] == target_dir_id] for f in dir_files: lines.append(f' {f["id"]} {f["filename"]} {_fmt_size(f)} {_fmt_date(f)}') @@ -620,9 +663,17 @@ async def _kb_grep(args: list[str], flags: set[str], user: dict, accessible = await _get_accessible_files(user, model_knowledge) if dir_scope: - # Scope to files under this directory path - accessible = [f for f in accessible if f.get('path', '').startswith(dir_scope + '/') or - (f.get('directory_id') and f.get('path', '').startswith(dir_scope))] + # Resolve directory and collect all descendant IDs + kb_ids = {fi['knowledge_id'] for fi in accessible if fi.get('knowledge_id')} + target_dir_ids = set() + for kb_id in kb_ids: + dir_id = await _resolve_dir_path(dir_scope, kb_id) + if dir_id: + desc = await _get_descendant_dir_ids(dir_id, kb_id) + target_dir_ids.update(desc) + if not target_dir_ids: + return f'No files found under "{dir_scope}/"' + accessible = [f for f in accessible if f.get('directory_id') in target_dir_ids] if not accessible: return f'No files found under "{dir_scope}/"' @@ -664,20 +715,20 @@ async def _kb_grep(args: list[str], flags: set[str], user: dict, for line_num, line_text in file_matches: if len(results) < MAX_GREP_MATCHES: results.append( - f'{file_info["id"]} {file_info.get("path", file_info["filename"])}:{line_num}: {line_text.rstrip()}' + f'{file_info["id"]} {file_info["filename"]}:{line_num}: {line_text.rstrip()}' ) if count_only: if not file_match_counts: return f'No matches for "{pattern}"' - lines = [f'{fi["id"]} {fi.get("path", fi["filename"])}: {cnt}' for fi, cnt in file_match_counts] + lines = [f'{fi["id"]} {fi["filename"]}: {cnt}' for fi, cnt in file_match_counts] lines.append(f'Total: {total_matches} matches in {len(file_match_counts)} files') return '\n'.join(lines) if filenames_only: if not files_with_matches: return f'No matches for "{pattern}"' - return '\n'.join(f'{fi["id"]} {fi.get("path", fi["filename"])}' for fi in files_with_matches) + return '\n'.join(f'{fi["id"]} {fi["filename"]}' for fi in files_with_matches) if not results: return f'No matches for "{pattern}" across {len(accessible)} files' @@ -707,8 +758,14 @@ async def _kb_find(args: list[str], flags: set[str], user: dict, accessible = await _get_accessible_files(user, model_knowledge) if dir_scope: - accessible = [f for f in accessible if f.get('path', '').startswith(dir_scope + '/') or - f.get('path', '').startswith(dir_scope)] + kb_ids = {fi['knowledge_id'] for fi in accessible if fi.get('knowledge_id')} + target_dir_ids = set() + for kb_id in kb_ids: + dir_id = await _resolve_dir_path(dir_scope, kb_id) + if dir_id: + desc = await _get_descendant_dir_ids(dir_id, kb_id) + target_dir_ids.update(desc) + accessible = [f for f in accessible if f.get('directory_id') in target_dir_ids] matched = [f for f in accessible if fnmatch.fnmatch(f['filename'], pattern)] @@ -719,7 +776,7 @@ async def _kb_find(args: list[str], flags: set[str], user: dict, lines = [] for f in matched: kb_info = f' ({f["knowledge_name"]})' if f.get('knowledge_name') else '' - lines.append(f'{f["id"]} {f.get("path", f["filename"])}{kb_info}') + lines.append(f'{f["id"]} {f["filename"]}{kb_info}') return '\n'.join(lines)