refactor(cli): Refactor clean command to use DB adapter layer

Author: Zhe Yu

src/vectorcode/database/base.py

@@ -123,7 +123,9 @@ async def delete(
         pass
 
     @abstractmethod
-    async def drop(self):
+    async def drop(
+        self, *, collection_id: str | None = None, collection_path: str | None = None
+    ):
         """
         Delete a collection (`self._configs.project_root`) from the database.
         """
@@ -196,3 +198,15 @@ async def get_chunks(self, file_path) -> list[Chunk]:
         If not found, return an empty list.
         """
         pass
+
+    async def cleanup(self) -> list[str]:
+        """
+        Remove empty collections from the database.
+        """
+        removed: list[str] = []
+        for collection in await self.list_collections():
+            if collection.chunk_count == 0:
+                removed.append(collection.path)
+                await self.drop(collection_path=collection.path)
+
+        return removed

src/vectorcode/database/chroma0.py

@@ -523,13 +523,12 @@ async def delete(self) -> int:
 
         return len(rm_paths)
 
-    async def drop(
-        self,
-    ):
-        collection_path = str(self._configs.project_root)
+    async def drop(self, *, collection_id=None, collection_path=None):
+        collection_id = collection_id or get_collection_id(
+            str(collection_path or self._configs.project_root)
+        )
         async with _Chroma0ClientManager().get_client(self._configs) as client:
-            await self._create_or_get_collection(collection_path, False)
-            await client.delete_collection(get_collection_id(collection_path))
+            await client.delete_collection(collection_id)
 
     async def get_chunks(self, file_path) -> list[Chunk]:
         file_path = os.path.abspath(file_path)

src/vectorcode/subcommands/clean.py

@@ -1,26 +1,16 @@
 import logging
-import os
-
-from chromadb.api import AsyncClientAPI
 
 from vectorcode.cli_utils import Config
-from vectorcode.common import ClientManager, get_collections
+from vectorcode.database import get_database_connector
 
 logger = logging.getLogger(name=__name__)
 
 
-async def run_clean_on_client(client: AsyncClientAPI, pipe_mode: bool):
-    async for collection in get_collections(client):
-        meta = collection.metadata
-        logger.debug(f"{meta.get('path')}: {await collection.count()} chunk(s)")
-        if await collection.count() == 0 or not os.path.isdir(meta["path"]):
-            await client.delete_collection(collection.name)
-            logger.info(f"Deleted collection for {meta['path']}")
-            if not pipe_mode:
-                print(f"Deleted {meta['path']}.")
-
-
 async def clean(configs: Config) -> int:
-    async with ClientManager().get_client(configs) as client:
-        await run_clean_on_client(client, configs.pipe)
-        return 0
+    database = get_database_connector(configs)
+    for removed in await database.cleanup():
+        message = f"Deleted collection: {removed}"
+        logger.info(message)
+        if not configs.pipe:
+            print(message)
+    return 0