refac

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 <path>, 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)