feat(cli): Refactor update command to use DB adapter layer

Author: Zhe Yu

src/vectorcode/subcommands/update.py

@@ -1,93 +1,65 @@
 import asyncio
 import logging
 import os
-import sys
 from asyncio import Lock
 
 import tqdm
-from chromadb.api.types import IncludeEnum
-from chromadb.errors import InvalidCollectionException
 
 from vectorcode.cli_utils import Config
-from vectorcode.common import ClientManager, get_collection, verify_ef
-from vectorcode.subcommands.vectorise import VectoriseStats, chunked_add, show_stats
+from vectorcode.database import get_database_connector
+from vectorcode.database.types import ResultType
+from vectorcode.database.utils import hash_file
+from vectorcode.subcommands.vectorise import (
+    VectoriseStats,
+    show_stats,
+    vectorise_worker,
+)
+from vectorcode.subcommands.vectorise.filter import FilterManager
 
 logger = logging.getLogger(name=__name__)
 
 
 async def update(configs: Config) -> int:
-    async with ClientManager().get_client(configs) as client:
-        try:
-            collection = await get_collection(client, configs, False)
-        except IndexError as e:
-            print(
-                f"{e.__class__.__name__}: Failed to get/create the collection. Please check your config."
-            )
-            return 1
-        except (ValueError, InvalidCollectionException) as e:
-            print(
-                f"{e.__class__.__name__}: There's no existing collection for {configs.project_root}",
-                file=sys.stderr,
-            )
-            return 1
-        if collection is None:  # pragma: nocover
-            logger.error(
-                f"Failed to find a collection at {configs.project_root} from {configs.db_url}"
-            )
-            return 1
-        if not verify_ef(collection, configs):  # pragma: nocover
-            return 1
+    assert configs.project_root is not None
+    database = get_database_connector(configs)
 
-        metas = (await collection.get(include=[IncludeEnum.metadatas]))["metadatas"]
-        if metas is None or len(metas) == 0:  # pragma: nocover
-            logger.debug("Empty collection.")
-            return 0
+    filters = FilterManager()
 
-        files_gen = (str(meta.get("path", "")) for meta in metas)
-        files = set()
-        orphanes = set()
-        for file in files_gen:
-            if os.path.isfile(file):
-                files.add(file)
-            else:
-                orphanes.add(file)
+    collection_files = (
+        await database.list_collection_content(what=ResultType.document)
+    ).files
 
-        stats = VectoriseStats(removed=len(orphanes))
-        collection_lock = Lock()
-        stats_lock = Lock()
-        max_batch_size = await client.get_max_batch_size()
-        semaphore = asyncio.Semaphore(os.cpu_count() or 1)
+    existing_hashes = set(i.sha256 for i in collection_files)
 
-        with tqdm.tqdm(
-            total=len(files), desc="Vectorising files...", disable=configs.pipe
-        ) as bar:
-            logger.info(f"Updating embeddings for {len(files)} file(s).")
-            try:
-                tasks = [
-                    asyncio.create_task(
-                        chunked_add(
-                            str(file),
-                            collection,
-                            collection_lock,
-                            stats,
-                            stats_lock,
-                            configs,
-                            max_batch_size,
-                            semaphore,
-                        )
-                    )
-                    for file in files
-                ]
-                for task in asyncio.as_completed(tasks):
-                    await task
-                    bar.update(1)
-            except asyncio.CancelledError:  # pragma: nocover
-                print("Abort.", file=sys.stderr)
-                return 1
+    files = (i.path for i in collection_files)
+    if not configs.force:
+        filters.add_filter(lambda x: hash_file(x) not in existing_hashes)
+    else:  # pragma: nocover
+        logger.info("Ignoring exclude specs.")
+
+    files = list(filters(files))
+    stats = VectoriseStats()
+    stats_lock = Lock()
+    semaphore = asyncio.Semaphore(os.cpu_count() or 1)
+
+    with tqdm.tqdm(
+        total=len(files), desc="Vectorising files...", disable=configs.pipe
+    ) as bar:
+        try:
+            tasks = [
+                asyncio.create_task(
+                    vectorise_worker(database, file, semaphore, stats, stats_lock)
+                )
+                for file in files
+            ]
+            for task in asyncio.as_completed(tasks):
+                await task
+                bar.update(1)
+        except asyncio.CancelledError:
+            logger.warning("Abort.")
+            return 1
 
-        if len(orphanes):
-            logger.info(f"Removing {len(orphanes)} orphaned files from database.")
-            await collection.delete(where={"path": {"$in": list(orphanes)}})
+    await database.check_orphanes()
 
-        show_stats(configs, stats)
-        return 0
+    show_stats(configs=configs, stats=stats)
+    return 0

src/vectorcode/subcommands/vectorise/__init__.py

@@ -3,7 +3,6 @@
 import hashlib
 import logging
 import os
-import sys
 import uuid
 from asyncio import Lock, Semaphore
 from typing import Iterable, Optional
@@ -306,7 +305,7 @@ async def vectorise(configs: Config) -> int:
                 await task
                 bar.update(1)
         except asyncio.CancelledError:
-            print("Abort.", file=sys.stderr)
+            logger.warning("Abort.")
             return 1
 
     await database.check_orphanes()

src/vectorcode/subcommands/vectorise/filter.py

@@ -36,11 +36,14 @@ def __call__(self, files: Iterable[str]) -> Iterable[str]:
                 f"Applying the following filters: {list(i.__name__ for i in self._filters)} to the following files ({len(files)}): {files}"
             )
 
-        for f in self._filters:
-            files = filter(f, files)
-
-            if self._has_debugging():  # pragma: nocover
-                files = tuple(files)
-                logger.debug(f"{f.__name__} remaining items ({len(files)}): {files}")
+        if self._filters:
+            for f in self._filters:
+                files = filter(f, files)
+
+                if self._has_debugging():  # pragma: nocover
+                    files = tuple(files)
+                    logger.debug(
+                        f"{f.__name__} remaining items ({len(files)}): {files}"
+                    )
 
         return files