refactor(cli): Separate embedding function from collection

Author: Zhe Yu

src/vectorcode/cli_utils.py

@@ -114,6 +114,9 @@ class Config:
     files_action: Optional[FilesAction] = None
     rm_paths: list[str] = field(default_factory=list)
 
+    def __hash__(self) -> int:
+        return hash(self.__repr__())
+
     @classmethod
     async def import_from(cls, config_dict: dict[str, Any]) -> "Config":
         """

src/vectorcode/common.py

@@ -8,6 +8,7 @@
 import sys
 from asyncio.subprocess import Process
 from dataclasses import dataclass
+from functools import cache
 from typing import Any, AsyncGenerator, Optional
 from urllib.parse import urlparse
 
@@ -127,11 +128,15 @@ def get_collection_name(full_path: str) -> str:
     return collection_id
 
 
-def get_embedding_function(configs: Config) -> chromadb.EmbeddingFunction | None:
+@cache
+def get_embedding_function(configs: Config) -> chromadb.EmbeddingFunction:
     try:
-        return getattr(embedding_functions, configs.embedding_function)(
+        ef = getattr(embedding_functions, configs.embedding_function)(
             **configs.embedding_params
         )
+        if ef is None:
+            raise AttributeError()
+        return ef
     except AttributeError:
         logger.warning(
             f"Failed to use {configs.embedding_function}. Falling back to Sentence Transformer.",
@@ -161,7 +166,6 @@ async def get_collection(
     full_path = str(expand_path(str(configs.project_root), absolute=True))
     if __COLLECTION_CACHE.get(full_path) is None:
         collection_name = get_collection_name(full_path)
-        embedding_function = get_embedding_function(configs)
 
         collection_meta: dict[str, str | int] = {
             "path": full_path,
@@ -183,14 +187,11 @@ async def get_collection(
             f"Getting/Creating collection with the following metadata: {collection_meta}"
         )
         if not make_if_missing:
-            __COLLECTION_CACHE[full_path] = await client.get_collection(
-                collection_name, embedding_function
-            )
+            __COLLECTION_CACHE[full_path] = await client.get_collection(collection_name)
         else:
             collection = await client.get_or_create_collection(
                 collection_name,
                 metadata=collection_meta,
-                embedding_function=embedding_function,
             )
             if (
                 not collection.metadata.get("hostname") == socket.gethostname()

src/vectorcode/subcommands/query/__init__.py

@@ -19,6 +19,7 @@
 from vectorcode.common import (
     ClientManager,
     get_collection,
+    get_embedding_function,
     verify_ef,
 )
 from vectorcode.subcommands.query.reranker import (
@@ -67,7 +68,7 @@ async def get_query_result_files(
                 )
                 logger.info(f"Querying {num_query} chunks for reranking.")
         results = await collection.query(
-            query_texts=query_chunks,
+            query_embeddings=get_embedding_function(configs)(query_chunks),
             n_results=num_query,
             include=[
                 IncludeEnum.metadatas,

src/vectorcode/subcommands/vectorise.py

@@ -28,6 +28,7 @@
 from vectorcode.common import (
     ClientManager,
     get_collection,
+    get_embedding_function,
     list_collection_files,
     verify_ef,
 )
@@ -92,6 +93,7 @@ async def chunked_add(
     max_batch_size: int,
     semaphore: asyncio.Semaphore,
 ):
+    embedding_function = get_embedding_function(configs)
     full_path_str = str(expand_path(str(file_path), True))
     orig_sha256 = None
     new_sha256 = hash_file(full_path_str)
@@ -147,6 +149,9 @@ async def chunked_add(
                     await collection.add(
                         ids=[get_uuid() for _ in inserted_chunks],
                         documents=[str(i) for i in inserted_chunks],
+                        embeddings=embedding_function(
+                            list(str(c) for c in inserted_chunks)
+                        ),
                         metadatas=metas,
                     )
     except (UnicodeDecodeError, UnicodeError):  # pragma: nocover