feat(cli): proof-of-concept query rewriter

Author: Davidyz

pyproject.toml

@@ -68,6 +68,7 @@ dev = [
     "pytest-asyncio>=0.25.3",
     "debugpy>=1.8.12",
     "coredumpy>=0.4.1",
+    "openai>=1.77.0",
 ]
 
 [project.optional-dependencies]

src/vectorcode/cli_utils.py

@@ -96,6 +96,8 @@ class Config:
     hnsw: dict[str, str | int] = field(default_factory=dict)
     chunk_filters: dict[str, list[str]] = field(default_factory=dict)
     encoding: str = "utf8"
+    rewriter: Optional[str] = None
+    rewriter_params: dict[str, Any] = field(default_factory=dict)
 
     @classmethod
     async def import_from(cls, config_dict: dict[str, Any]) -> "Config":
@@ -145,6 +147,10 @@ async def import_from(cls, config_dict: dict[str, Any]) -> "Config":
                     "chunk_filters", default_config.chunk_filters
                 ),
                 "encoding": config_dict.get("encoding", default_config.encoding),
+                "rewriter": config_dict.get("rewriter"),
+                "rewriter_params": config_dict.get(
+                    "rewriter_params", default_config.rewriter_params
+                ),
             }
         )
 

src/vectorcode/rewriter/__init__.py

@@ -0,0 +1,4 @@
+from .base import RewriterBase
+from .openai import OpenAIRewriter
+
+__all__ = ["RewriterBase", "OpenAIRewriter"]

src/vectorcode/rewriter/base.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+
+from vectorcode.cli_utils import Config
+
+
+class RewriterBase(ABC):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        self.config = config
+
+    @abstractmethod
+    async def rewrite(self, original_query: list[str]) -> list[str]:
+        raise NotImplementedError

src/vectorcode/rewriter/openai.py

@@ -0,0 +1,74 @@
+import logging
+from typing import override
+
+import openai
+from openai.types.chat import ChatCompletion
+from pydantic import BaseModel, Field
+
+from vectorcode.cli_utils import Config
+from vectorcode.rewriter.base import RewriterBase
+
+logger = logging.getLogger(name=__name__)
+
+
+class _NewQuery(BaseModel):
+    keywords: list[str] = Field(
+        description="Orthogonal keywords for the vector search."
+    )
+
+
+class OpenAIRewriter(RewriterBase):
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+        self.client = openai.Client(
+            **self.config.rewriter_params.get("client_kwargs", {})
+        )
+        self.system_prompt = """
+Role:
+        You are a code-aware rewriter that improves technical queries/docs for retrieval. Never assume a programming language unless the input explicitly includes syntax, APIs, or error messages from one.
+Rules:
+
+    For Queries:
+
+        Fix unambiguous typos (e.g., "Pytoch" → "PyTorch").
+        
+        Omit langauge-specific keywords (e.g., "async def foo():" → "foo")
+
+    For Docs/Code:
+
+        Clarify ambiguous terms only with explicit context.
+
+        Never modify code logic or variable names.
+
+Anti-Goals:
+
+    No language assumptions.
+
+    No code changes.
+
+    No hallucinations.
+"""
+
+    @override
+    async def rewrite(self, original_query: list[str]):
+        comp: ChatCompletion = self.client.beta.chat.completions.parse(
+            messages=[
+                {"role": "system", "content": self.system_prompt},
+                {"role": "user", "content": " ".join(original_query)},
+            ],
+            response_format=_NewQuery,
+            **self.config.rewriter_params.get("completion_kwargs", {}),
+        )
+        if comp is None or len(comp.choices) == 0:
+            logger.info("Recieved no rewritten query. Fallingback to original_query.")
+            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:
+            logger.warning(
+                f"Failed to parse structured output: {choice.refusal}. Fallingback to original_query."
+            )
+            return original_query

src/vectorcode/subcommands/query/__init__.py

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