feat(db): Refactor database connectors to use embeddings and improve listing

Author: Zhe Yu

src/vectorcode/database/base.py

@@ -9,6 +9,7 @@
 
 from vectorcode.chunking import TreeSitterChunker
 from vectorcode.cli_utils import Config
+from vectorcode.common import get_embedding_function
 from vectorcode.database.types import (
     CollectionContent,
     CollectionInfo,
@@ -64,7 +65,7 @@ async def count(self, what: ResultType = ResultType.chunk) -> int:
         """
         Returns the chunk count or file count of the given collection, depending on the value passed for `what`.
         """
-        collection_content = await self.list(what)
+        collection_content = await self.list_collection_content(what)
         match what:
             case ResultType.chunk:
                 return len(collection_content.chunks)
@@ -74,7 +75,6 @@ async def count(self, what: ResultType = ResultType.chunk) -> int:
     @abstractmethod
     async def query(
         self,
-        keywords_embeddings: list[NDArray],
     ) -> Sequence[QueryResult]:
         pass
 
@@ -99,7 +99,9 @@ async def list_collections(self) -> Sequence[CollectionInfo]:
         pass
 
     @abstractmethod
-    async def list(self, what: Optional[ResultType] = None) -> CollectionContent:
+    async def list_collection_content(
+        self, what: Optional[ResultType] = None
+    ) -> CollectionContent:
         """
         List the content of a collection (from `self._configs.project_root`).
 
@@ -163,7 +165,7 @@ async def check_orphanes(self) -> int:
         """
 
         orphanes: list[str] = []
-        database_files = (await self.list(ResultType.document)).files
+        database_files = (await self.list_collection_content(ResultType.document)).files
         for file in database_files:
             path = file.path
             if not os.path.isfile(path):
@@ -174,3 +176,14 @@ async def check_orphanes(self) -> int:
         await self.delete()
 
         return len(orphanes)
+
+    def get_embedding(self, texts: str | list[str]) -> list[NDArray]:
+        """
+        Generate embeddings and truncate them to `self._configs.embedding_dims` if needed.
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+        embeddings = get_embedding_function(self._configs)(texts)
+        if self._configs.embedding_dims:
+            embeddings = [e[: self._configs.embedding_dims] for e in embeddings]
+        return embeddings

src/vectorcode/database/chroma0.py

@@ -267,12 +267,14 @@ def __init__(self, configs: Config):
         params.update(self._configs.db_params)
         self._configs.db_params = params
 
-    async def query(self, keywords_embeddings):
+    async def query(self):
         assert self._configs.query is not None
         assert len(self._configs.query), "Keywords cannot be empty"
+        keywords_embeddings = self.get_embedding(self._configs.query)
         assert len(keywords_embeddings) == len(self._configs.query), (
             "Number of embeddings must match number of keywords."
         )
+
         collection_path = str(self._configs.project_root)
         collection: AsyncCollection = await self._create_or_get_collection(
             collection_path=collection_path, allow_create=False

src/vectorcode/subcommands/files/ls.py

@@ -9,7 +9,7 @@
 
 async def ls(configs: Config) -> int:
     database = get_database_connector(configs)
-    files = list(i.path for i in (await database.list()).files)
+    files = list(i.path for i in (await database.list_collection_content()).files)
     if configs.pipe:
         print(json.dumps(files))
     else: