refactor(cli): Stop hardcoding rewriter

Author: Zhe Yu

src/vectorcode/cli_utils.py

@@ -97,6 +97,7 @@ class Config:
     filetype_map: dict[str, list[str]] = field(default_factory=dict)
     encoding: str = "utf8"
     hooks: bool = False
+    use_rewriter: bool = False
     rewriter: Optional[str] = None
     rewriter_params: dict[str, Any] = field(default_factory=dict)
 
@@ -317,6 +318,11 @@ def get_cli_parser():
         help="What to include in the final output.",
         default=__default_config.include,
     )
+    query_parser.add_argument(
+        "--rewrite",
+        action="store_true",
+        help="Apply rewriter to rewrite the query before running the search.",
+    )
 
     subparsers.add_parser("drop", parents=[shared_parser], help="Remove a collection.")
     hooks_parser = subparsers.add_parser(
@@ -422,6 +428,7 @@ async def parse_cli_args(args: Optional[Sequence[str]] = None):
             configs_items["use_absolute_path"] = main_args.absolute
             configs_items["include"] = [QueryInclude(i) for i in main_args.include]
             configs_items["encoding"] = main_args.encoding
+            configs_items["use_rewriter"] = main_args.rewrite
         case "check":
             configs_items["check_item"] = main_args.check_item
         case "init":

src/vectorcode/rewriter/__init__.py

@@ -1,4 +1,28 @@
+import logging
+import sys
+from typing import Optional
+
+from vectorcode.cli_utils import Config
+
 from .base import RewriterBase
 from .openai import OpenAIRewriter
 
+logger = logging.getLogger(name=__name__)
 __all__ = ["RewriterBase", "OpenAIRewriter"]
+
+
+class RewriterError(Exception):
+    pass
+
+
+def get_rewriter(configs: Config) -> Optional[RewriterBase]:
+    if configs.rewriter is None:
+        logger.warning("Rewriter hasn't been configured. Skipping rewriting.")
+        return None
+    if configs.rewriter == "RewriterBase":
+        raise RewriterError("RewriterBase is not a valid rewriter!")
+    rewriter_cls = getattr(sys.modules[__name__], configs.rewriter)
+    if issubclass(rewriter_cls, RewriterBase):
+        logger.info(f"Loaded {configs.rewriter}")
+        return rewriter_cls(configs)
+    raise RewriterError(f"Failed to find {configs.rewriter}!")

src/vectorcode/rewriter/openai.py

@@ -64,7 +64,6 @@ async def rewrite(self, original_query: list[str]):
             return original_query
         choice = comp.choices[0].message
         if choice and choice.parsed:
-            print(choice.parsed)
             logger.debug(f"Rewritten queries to: {choice.parsed}")
             return choice.parsed.keywords
         else:

src/vectorcode/subcommands/query/__init__.py

@@ -20,7 +20,6 @@
     get_collection,
     verify_ef,
 )
-from vectorcode.rewriter import OpenAIRewriter
 from vectorcode.subcommands.query.reranker import (
     RerankerError,
     get_reranker,
@@ -34,8 +33,12 @@ async def get_query_result_files(
 ) -> list[str]:
     query_chunks = []
     if configs.query:
-        if configs.rewriter:
-            configs.query = await OpenAIRewriter(configs).rewrite(configs.query)
+        if configs.use_rewriter:
+            from vectorcode.rewriter import get_rewriter  # lazy import
+
+            rewriter = get_rewriter(configs)
+            if rewriter is not None:
+                configs.query = await rewriter.rewrite(configs.query)
         chunker = StringChunker(configs)
         for q in configs.query:
             query_chunks.extend(str(i) for i in chunker.chunk(q))