feat(cli): add files subcommand to the LSP server

Author: Zhe Yu

src/vectorcode/common.py

@@ -229,13 +229,15 @@ def verify_ef(collection: AsyncCollection, configs: Config):
 
 
 async def list_collection_files(collection: AsyncCollection) -> list[str]:
-    return list(
-        set(
-            str(c.get("path", None))
-            for c in (await collection.get(include=[IncludeEnum.metadatas])).get(
-                "metadatas"
+    return sorted(
+        list(
+            set(
+                str(c.get("path", None))
+                for c in (await collection.get(include=[IncludeEnum.metadatas])).get(
+                    "metadatas"
+                )
+                or []
             )
-            or []
         )
     )
 

src/vectorcode/lsp_main.py

@@ -6,8 +6,10 @@
 import time
 import traceback
 import uuid
+from typing import cast
 
 import shtab
+from chromadb.types import Where
 
 from vectorcode.subcommands.vectorise import (
     VectoriseStats,
@@ -35,14 +37,16 @@
 from vectorcode import __version__
 from vectorcode.cli_utils import (
     CliAction,
+    FilesAction,
     cleanup_path,
     config_logging,
     expand_globs,
+    expand_path,
     find_project_root,
     get_project_config,
     parse_cli_args,
 )
-from vectorcode.common import ClientManager, get_collection
+from vectorcode.common import ClientManager, get_collection, list_collection_files
 from vectorcode.subcommands.ls import get_collection_list
 from vectorcode.subcommands.query import build_query_results
 
@@ -105,7 +109,11 @@ async def execute_command(ls: LanguageServer, args: list[str]):
         async with ClientManager().get_client(final_configs) as client:
             progress_token = str(uuid.uuid4())
 
-            if final_configs.action in {CliAction.vectorise, CliAction.query}:
+            if final_configs.action in {
+                CliAction.vectorise,
+                CliAction.query,
+                CliAction.files,
+            }:
                 collection = await get_collection(
                     client=client,
                     configs=final_configs,
@@ -222,6 +230,40 @@ async def execute_command(ls: LanguageServer, args: list[str]):
                         ),
                     )
                     return stats.to_dict()
+                case CliAction.files:
+                    if collection is None:
+                        return
+                    match final_configs.files_action:
+                        case FilesAction.ls:
+                            return await list_collection_files(collection)
+                        case FilesAction.rm:
+                            to_be_removed = list(
+                                str(expand_path(p, True))
+                                for p in final_configs.rm_paths
+                                if os.path.isfile(p)
+                            )
+                            if len(to_be_removed) == 0:
+                                return
+                            ls.progress.begin(
+                                progress_token,
+                                types.WorkDoneProgressBegin(
+                                    title="VectorCode",
+                                    message=f"Removing {len(to_be_removed)} file(s).",
+                                ),
+                            )
+                            await collection.delete(
+                                where=cast(
+                                    Where,
+                                    {"path": {"$in": to_be_removed}},
+                                )
+                            )
+                            ls.progress.begin(
+                                progress_token,
+                                types.WorkDoneProgressBegin(
+                                    title="VectorCode",
+                                    message="Removal finished.",
+                                ),
+                            )
                 case _ as c:  # pragma: nocover
                     error_message = f"Unsupported vectorcode subcommand: {str(c)}"
                     logger.error(

src/vectorcode/subcommands/files/ls.py

@@ -1,33 +1,20 @@
 import json
 import logging
 
-from chromadb.api.models.AsyncCollection import AsyncCollection
-from chromadb.api.types import IncludeEnum
-
 from vectorcode.cli_utils import Config
-from vectorcode.common import ClientManager, get_collection
+from vectorcode.common import ClientManager, get_collection, list_collection_files
 
 logger = logging.getLogger(name=__name__)
 
 
-async def list_files(collection: AsyncCollection) -> list[str]:
-    meta = (await collection.get(include=[IncludeEnum.metadatas])).get("metadatas")
-    if meta is None:
-        logger.warning("Failed to fetch metadatas from the database.")
-        return []
-    paths: list[str] = list(set(str(m.get("path")) for m in meta))
-    paths.sort()
-    return paths
-
-
 async def ls(configs: Config) -> int:
     async with ClientManager().get_client(configs=configs) as client:
         try:
             collection = await get_collection(client, configs, False)
         except ValueError:
             logger.error(f"There's no existing collection at {configs.project_root}.")
             return 1
-        paths = await list_files(collection)
+        paths = await list_collection_files(collection)
         if configs.pipe:
             print(json.dumps(list(paths)))
         else:

tests/test_lsp.py

@@ -1,3 +1,4 @@
+import os
 from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -6,7 +7,7 @@
 from pygls.server import LanguageServer
 
 from vectorcode import __version__
-from vectorcode.cli_utils import CliAction, Config, QueryInclude
+from vectorcode.cli_utils import CliAction, Config, FilesAction, QueryInclude
 from vectorcode.lsp_main import (
     execute_command,
     lsp_start,
@@ -374,3 +375,124 @@ async def test_execute_command_no_default_project_root(
         with pytest.raises((AssertionError, JsonRpcInternalError)):
             await execute_command(mock_language_server, ["query", "test"])
     DEFAULT_PROJECT_ROOT = None  # Reset the global variable
+
+
+@pytest.mark.asyncio
+async def test_execute_command_files_ls(mock_language_server, mock_config: Config):
+    mock_config.action = CliAction.files
+    mock_config.files_action = FilesAction.ls
+    mock_config.project_root = "/test/project"
+
+    dummy_files = ["/test/project/file1.py", "/test/project/file2.txt"]
+    mock_client = AsyncMock()
+    mock_collection = AsyncMock()
+
+    with (
+        patch(
+            "vectorcode.lsp_main.parse_cli_args", new_callable=AsyncMock
+        ) as mock_parse_cli_args,
+        patch(
+            "vectorcode.lsp_main.ClientManager._create_client", return_value=mock_client
+        ),
+        patch("vectorcode.lsp_main.get_collection", return_value=mock_collection),
+        patch(
+            "vectorcode.lsp_main.list_collection_files", return_value=dummy_files
+        ) as mock_list_collection_files,
+    ):
+        mock_parse_cli_args.return_value = mock_config
+
+        mock_config.merge_from = AsyncMock(return_value=mock_config)
+
+        result = await execute_command(mock_language_server, ["files", "ls"])
+
+        assert result == dummy_files
+        mock_language_server.progress.create_async.assert_called_once()
+
+        mock_list_collection_files.assert_called_once_with(mock_collection)
+        # For 'ls' action, progress.begin/end are not explicitly called in the lsp_main,
+        # but create_async is called before the match statement.
+        mock_language_server.progress.begin.assert_not_called()
+        mock_language_server.progress.end.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_execute_command_files_rm(mock_language_server, mock_config: Config):
+    mock_config.action = CliAction.files
+    mock_config.files_action = FilesAction.rm
+    mock_config.project_root = "/test/project"
+    mock_config.rm_paths = ["file_to_remove.py", "another_file.txt"]
+
+    expanded_paths = [
+        "/test/project/file_to_remove.py",
+        "/test/project/another_file.txt",
+    ]
+    mock_client = AsyncMock()
+    mock_collection = AsyncMock()
+
+    with (
+        patch(
+            "vectorcode.lsp_main.parse_cli_args", new_callable=AsyncMock
+        ) as mock_parse_cli_args,
+        patch(
+            "vectorcode.lsp_main.ClientManager._create_client", return_value=mock_client
+        ),
+        patch("vectorcode.lsp_main.get_collection", return_value=mock_collection),
+        patch(
+            "os.path.isfile",
+            side_effect=lambda x: x in expanded_paths or x in mock_config.rm_paths,
+        ),
+        patch(
+            "vectorcode.lsp_main.expand_path",
+            side_effect=lambda p, *args: os.path.join(mock_config.project_root, p),
+        ),
+    ):
+        mock_parse_cli_args.return_value = mock_config
+
+        mock_config.merge_from = AsyncMock(return_value=mock_config)
+
+        await execute_command(
+            mock_language_server,
+            ["files", "rm", "file_to_remove.py", "another_file.txt"],
+        )
+
+        mock_collection.delete.assert_called_once_with(
+            where={"path": {"$in": expanded_paths}}
+        )
+
+
+@pytest.mark.asyncio
+async def test_execute_command_files_rm_no_files_to_remove(
+    mock_language_server, mock_config: Config
+):
+    mock_config.action = CliAction.files
+    mock_config.files_action = FilesAction.rm
+    mock_config.project_root = "/test/project"
+    mock_config.rm_paths = ["non_existent_file.py"]
+
+    mock_client = AsyncMock()
+    mock_collection = AsyncMock()
+
+    with (
+        patch(
+            "vectorcode.lsp_main.parse_cli_args", new_callable=AsyncMock
+        ) as mock_parse_cli_args,
+        patch(
+            "vectorcode.lsp_main.ClientManager._create_client", return_value=mock_client
+        ),
+        patch("vectorcode.lsp_main.get_collection", return_value=mock_collection),
+        patch("os.path.isfile", return_value=False),
+        patch(
+            "vectorcode.lsp_main.expand_path",
+            side_effect=lambda p, *args: os.path.join(mock_config.project_root, p),
+        ),
+    ):
+        mock_parse_cli_args.return_value = mock_config
+
+        mock_config.merge_from = AsyncMock(return_value=mock_config)
+
+        result = await execute_command(
+            mock_language_server, ["files", "rm", "non_existent_file.py"]
+        )
+
+        assert result is None
+        mock_collection.delete.assert_not_called()