refactor(cli): ls in CLI mode now use the DB adapter layer.

Author: Davidyz

src/vectorcode/database/__init__.py

@@ -16,6 +16,9 @@ def get_database_connector(config: Config) -> DatabaseConnectorBase:
     """
     cls: Optional[Type[DatabaseConnectorBase]] = None
 
+    if not config.db_type.endswith("Connector"):
+        config.db_type = f"{config.db_type}Connector"
+
     match config.db_type:
         case "ChromaDB0Connector":
             from vectorcode.database.chroma0 import ChromaDB0Connector

src/vectorcode/database/chroma0.py

@@ -20,7 +20,6 @@
 from chromadb.errors import InvalidCollectionException
 from tree_sitter import Point
 
-import vectorcode.subcommands.query.types as vectorcode_query_types
 from vectorcode.chunking import Chunk, TreeSitterChunker
 from vectorcode.cli_utils import Config, LockManager, expand_path
 from vectorcode.common import get_embedding_function
@@ -34,21 +33,22 @@
     VectoriseStats,
 )
 from vectorcode.database.utils import get_collection_id, hash_file
+from vectorcode.subcommands.query import types
 from vectorcode.subcommands.vectorise import get_uuid
 
 _logger = logging.getLogger(name=__name__)
 
 
 def __convert_chroma_query_results(
     chroma_result: QueryResult, queries: list[str]
-) -> list[vectorcode_query_types.QueryResult]:
+) -> list[types.QueryResult]:
     """Convert chromadb query result to in-house query results"""
     assert chroma_result["documents"] is not None
     assert chroma_result["distances"] is not None
     assert chroma_result["metadatas"] is not None
     assert chroma_result["ids"] is not None
 
-    chroma_results_list: list[vectorcode_query_types.QueryResult] = []
+    chroma_results_list: list[types.QueryResult] = []
     for q_i in range(len(queries)):
         q = queries[q_i]
         documents = chroma_result["documents"][q_i]
@@ -64,7 +64,7 @@ def __convert_chroma_query_results(
             if meta.get("path"):
                 chunk.path = str(meta["path"])
             chroma_results_list.append(
-                vectorcode_query_types.QueryResult(
+                types.QueryResult(
                     chunk=chunk,
                     path=str(meta.get("path", "")),
                     query=(q,),
@@ -391,11 +391,12 @@ async def list_collections(self):
             result: list[CollectionInfo] = []
             for col_name in await client.list_collections():
                 col = await client.get_collection(col_name)
-                col_counts = await self.list(col_name)
+                project_root = str(col.metadata.get("path"))
+                col_counts = await self.list(project_root)
                 result.append(
                     CollectionInfo(
                         id=col_name,
-                        path=str(col.metadata.get("path")),
+                        path=project_root,
                         embedding_function=col.metadata.get(
                             "embedding_function",
                             Config().embedding_function,  # fallback to default
@@ -413,9 +414,7 @@ async def list(self, collection_path, what=None) -> CollectionContent:
         Otherwise, this method may populate only one of them to save waiting time.
         """
         content = CollectionContent()
-        collection = await self._create_or_get_collection(
-            get_collection_id(collection_path)
-        )
+        collection = await self._create_or_get_collection((collection_path))
         raw_content = await collection.get(
             include=[
                 IncludeEnum.metadatas,

src/vectorcode/database/types.py

@@ -67,6 +67,15 @@ class CollectionInfo:
     chunk_count: int = 0
     metadata: dict[str, Any] = field(default_factory=dict)
 
+    def to_dict(self) -> dict[str, int | str]:
+        return {
+            "project-root": self.path,
+            "size": self.chunk_count,
+            "num_files": self.file_count,
+            "collection_name": self.id,
+            "embedding_function": self.embedding_function,
+        }
+
 
 @dataclass
 class FileInCollection:

src/vectorcode/subcommands/ls.py

@@ -8,7 +8,8 @@
 from chromadb.api.types import IncludeEnum
 
 from vectorcode.cli_utils import Config, cleanup_path
-from vectorcode.common import ClientManager, get_collections
+from vectorcode.common import get_collections
+from vectorcode.database import get_database_connector
 
 logger = logging.getLogger(name=__name__)
 
@@ -36,34 +37,34 @@ async def get_collection_list(client: AsyncClientAPI) -> list[dict]:
 
 
 async def ls(configs: Config) -> int:
-    async with ClientManager().get_client(configs) as client:
-        result: list[dict] = await get_collection_list(client)
-        logger.info(f"Found the following collections: {result}")
+    result = [
+        i.to_dict() for i in await get_database_connector(configs).list_collections()
+    ]
 
-        if configs.pipe:
-            print(json.dumps(result))
-        else:
-            table = []
-            for meta in result:
-                project_root = meta["project-root"]
-                if os.environ.get("HOME"):
-                    project_root = project_root.replace(os.environ["HOME"], "~")
-                row = [
-                    project_root,
-                    meta["size"],
-                    meta["num_files"],
-                    meta["embedding_function"],
-                ]
-                table.append(row)
-            print(
-                tabulate.tabulate(
-                    table,
-                    headers=[
-                        "Project Root",
-                        "Collection Size",
-                        "Number of Files",
-                        "Embedding Function",
-                    ],
-                )
+    if configs.pipe:
+        print(json.dumps(result))
+    else:
+        table = []
+        for meta in result:
+            project_root = str(meta["project-root"])
+            if os.environ.get("HOME"):
+                project_root = project_root.replace(os.environ["HOME"], "~")
+            row = [
+                project_root,
+                meta["size"],
+                meta["num_files"],
+                meta["embedding_function"],
+            ]
+            table.append(row)
+        print(
+            tabulate.tabulate(
+                table,
+                headers=[
+                    "Project Root",
+                    "Number of Embeddings",
+                    "Number of Files",
+                    "Embedding Function",
+                ],
             )
-        return 0
+        )
+    return 0