Add an LLM retriever (#70)

README.md

@@ -225,10 +225,16 @@ Currently, we support the following types of retrieval:
     - Note this is not available when running locally, only when using Pinecone as a vector store.
     - Contrary to [Anthropic's findings](https://www.anthropic.com/news/contextual-retrieval), we find that BM25 is actually damaging performance *on codebases*, because it gives undeserved advantage to Markdown files.
 
-- **Multi-query retrieval** performs multiple query rewrites, makes a separate retrieval call for each, and takes the union of the retrieved documents. You can activate it by passing `--multi-query-retrieval`.
+- **Multi-query retrieval** performs multiple query rewrites, makes a separate retrieval call for each, and takes the union of the retrieved documents. You can activate it by passing `--multi-query-retrieval`. This can be combined with both vanilla and hybrid RAG.
 
     - We find that [on our benchmark](benchmark/retrieval/README.md) this only marginally improves retrieval quality (from 0.44 to 0.46 R-precision) while being significantly slower and more expensive due to LLM calls. But your mileage may vary.
 
+- **LLM-only retrieval** completely circumvents indexing the codebase. We simply enumerate all file paths and pass them to an LLM together with the user query. We ask the LLM which files are likely to be relevant for the user query, solely based on their filenames. You can activate it by passing `--llm-retriever`.
+
+    - We find that [on our benchmark](benchmark/retrieval/README.md) the performance is comparable with vector database solutions (R-precision is 0.44 for both). This is quite remarkable, since we've saved so much effort by not indexing the codebase. However, we are reluctant to claim that these findings generalize, for the following reasons:
+        - Our (artificial) dataset occasionally contains explicit path names in the query, making it trivial for the LLM. Sample query: *"Alice is managing a series of machine learning experiments. Please explain in detail how `main` in `examples/pytorch/image-pretraining/run_mim.py` allows her to organize the outputs of each experiment in separate directories."*
+        - Our benchmark focuses on the Transformers library, which is well-maintained and the file paths are often meaningful. This might not be the case for all codebases.
+
 </details>
 
 # Why chat with a codebase?

benchmarks/retrieval/retrieve.py

@@ -41,24 +41,12 @@ def main():
     )
 
     parser.add("--max-instances", default=None, type=int, help="Maximum number of instances to process.")
-    sage.config.add_config_args(parser)
-    sage.config.add_llm_args(parser)  # Needed for --multi-query-retriever, which rewrites the query with an LLM.
-    sage.config.add_embedding_args(parser)
-    sage.config.add_vector_store_args(parser)
-    sage.config.add_reranking_args(parser)
-    sage.config.add_repo_args(parser)
-    sage.config.add_indexing_args(parser)
+
+    validator = sage.config.add_all_args(parser)
     args = parser.parse_args()
-    sage.config.validate_vector_store_args(args)
-    repo_manager = GitHubRepoManager(
-        args.repo_id,
-        commit_hash=args.commit_hash,
-        access_token=os.getenv("GITHUB_TOKEN"),
-        local_dir=args.local_dir,
-        inclusion_file=args.include,
-        exclusion_file=args.exclude,
-    )
-    repo_manager.download()
+    validator(args)
+
+    repo_manager = GitHubRepoManager.from_args(args)
     retriever = build_retriever_from_args(args, repo_manager)
 
     with open(args.benchmark, "r") as f:

pyproject.toml

@@ -23,6 +23,7 @@ classifiers = [
 dependencies = [
     "GitPython==3.1.43",
     "Pygments==2.18.0",
+    "anytree==2.12.1",
     "cohere==5.9.2",
     "configargparse",
     "fastapi==0.112.2",
@@ -46,6 +47,7 @@ dependencies = [
     "pinecone==5.0.1",
     "pinecone-text==0.9.0",
     "python-dotenv==1.0.1",
+    "python-Levenshtein==0.26.0",
     "requests==2.32.3",
     "semchunk==2.2.0",
     "sentence-transformers==3.1.0",

sage/chat.py

@@ -71,20 +71,10 @@ def main():
         default=False,
         help="Whether to make the gradio app publicly accessible.",
     )
-    sage_config.add_config_args(parser)
-
-    arg_validators = [
-        sage_config.add_repo_args(parser),
-        sage_config.add_embedding_args(parser),
-        sage_config.add_vector_store_args(parser),
-        sage_config.add_reranking_args(parser),
-        sage_config.add_llm_args(parser),
-    ]
 
+    validator = sage_config.add_all_args(parser)
     args = parser.parse_args()
-
-    for validator in arg_validators:
-        validator(args)
+    validator(args)
 
     rag_chain = build_rag_chain(args)
 

sage/config.py

@@ -71,6 +71,7 @@ def add_config_args(parser: ArgumentParser):
     args, _ = parser.parse_known_args()
     config_file = pkg_resources.resource_filename(__name__, f"configs/{args.mode}.yaml")
     parser.set_defaults(config=config_file)
+    return lambda _: True
 
 
 def add_repo_args(parser: ArgumentParser) -> Callable:
@@ -157,6 +158,14 @@ def add_vector_store_args(parser: ArgumentParser) -> Callable:
         help="When set to True, we rewrite the query 5 times, perform retrieval for each rewrite, and take the union "
         "of retrieved documents. See https://python.langchain.com/v0.1/docs/modules/data_connection/retrievers/MultiQueryRetriever/.",
     )
+    parser.add(
+        "--llm-retriever",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="When set to True, we use an LLM for retrieval: we pass the repository file hierarchy together with the "
+        "user query and ask the LLM to choose relevant files solely based on their paths. No indexing will be done, so "
+        "all the vector store / embedding arguments will be ignored.",
+    )
     return validate_vector_store_args
 
 
@@ -221,6 +230,25 @@ def add_llm_args(parser: ArgumentParser) -> Callable:
     return lambda _: True
 
 
+def add_all_args(parser: ArgumentParser) -> Callable:
+    """Adds all arguments to the parser and returns a validator."""
+    arg_validators = [
+        add_config_args(parser),
+        add_repo_args(parser),
+        add_embedding_args(parser),
+        add_vector_store_args(parser),
+        add_reranking_args(parser),
+        add_indexing_args(parser),
+        add_llm_args(parser),
+    ]
+
+    def validate_all(args):
+        for validator in arg_validators:
+            validator(args)
+
+    return validate_all
+
+
 def validate_repo_args(args):
     """Validates the configuration of the repository."""
     if not re.match(r"^[^/]+/[^/]+$", args.repo_id):
@@ -233,7 +261,7 @@ def _validate_openai_embedding_args(args):
         raise ValueError("Please set the OPENAI_API_KEY environment variable.")
 
     if not args.embedding_model:
-        args.embedding_model = "text-embedding-ada-002"
+        args.embedding_model = "text-embedding-3-small"
 
     if args.embedding_model not in OPENAI_DEFAULT_EMBEDDING_SIZE.keys():
         raise ValueError(f"Unrecognized embeddings.model={args.embedding_model}")
@@ -347,6 +375,19 @@ def validate_embedding_args(args):
 
 def validate_vector_store_args(args):
     """Validates the configuration of the vector store and sets defaults."""
+    if args.llm_retriever:
+        if not os.getenv("ANTHROPIC_API_KEY"):
+            raise ValueError(
+                "Please set the ANTHROPIC_API_KEY environment variable to use the LLM retriever. "
+                "(We're constrained to Claude because we need prompt caching.)"
+            )
+
+        if args.index_issues:
+            # The LLM retriever only makes sense on the code repository, since it passes file paths to the LLM.
+            raise ValueError("Cannot use --index-issues with --llm-retriever.")
+
+        # When using an LLM retriever, all the vector store arguments are ignored.
+        return
 
     if not args.index_namespace:
         # Attempt to derive a default index namespace from the repository information.

sage/data_manager.py

@@ -175,15 +175,14 @@ class GitHubRepoManager(DataManager):
             )
         return True
 
-    def walk(self) -> Generator[Tuple[Any, Dict], None, None]:
+    def walk(self, get_content: bool = True) -> Generator[Tuple[Any, Dict], None, None]:
         """Walks the local repository path and yields a tuple of (content, metadata) for each file.
         The filepath is relative to the root of the repository (e.g. "org/repo/your/file/path.py").
 
         Args:
-            included_extensions: Optional set of extensions to include.
-            excluded_extensions: Optional set of extensions to exclude.
+            get_content: When set to True, yields (content, metadata) tuples. When set to False, yields metadata only.
         """
-        # We will keep apending to these files during the iteration, so we need to clear them first.
+        # We will keep appending to these files during the iteration, so we need to clear them first.
         repo_name = self.repo_id.replace("/", "_")
         included_log_file = os.path.join(self.log_dir, f"included_{repo_name}.txt")
         excluded_log_file = os.path.join(self.log_dir, f"excluded_{repo_name}.txt")
@@ -208,20 +207,49 @@ class GitHubRepoManager(DataManager):
                     f.write(path + "\n")
 
             for file_path in included_file_paths:
+                relative_file_path = file_path[len(self.local_dir) + 1 :]
+                metadata = {
+                    "file_path": relative_file_path,
+                    "url": self.url_for_file(relative_file_path),
+                }
+
+                if not get_content:
+                    yield metadata
+                    continue
+
                 with open(file_path, "r") as f:
                     try:
                         contents = f.read()
                     except UnicodeDecodeError:
                         logging.warning("Unable to decode file %s. Skipping.", file_path)
                         continue
-                    relative_file_path = file_path[len(self.local_dir) + 1 :]
-                    metadata = {
-                        "file_path": relative_file_path,
-                        "url": self.url_for_file(relative_file_path),
-                    }
                     yield contents, metadata
 
     def url_for_file(self, file_path: str) -> str:
         """Converts a repository file path to a GitHub link."""
         file_path = file_path[len(self.repo_id) + 1 :]
         return f"https://github.com/{self.repo_id}/blob/{self.default_branch}/{file_path}"
+
+    def read_file(self, relative_file_path: str) -> str:
+        """Reads the content of the file at the given path."""
+        file_path = os.path.join(self.local_dir, relative_file_path)
+        with open(file_path, "r") as f:
+            return f.read()
+
+    def from_args(args: Dict):
+        """Creates a GitHubRepoManager from command-line arguments and clones the underlying repository."""
+        repo_manager = GitHubRepoManager(
+            repo_id=args.repo_id,
+            commit_hash=args.commit_hash,
+            access_token=os.getenv("GITHUB_TOKEN"),
+            local_dir=args.local_dir,
+            inclusion_file=args.include,
+            exclusion_file=args.exclude,
+        )
+        success = repo_manager.download()
+        if not success:
+            raise ValueError(
+                f"Unable to clone {args.repo_id}. Please check that it exists and you have access to it. "
+                "For private repositories, please set the GITHUB_TOKEN variable in your environment."
+            )
+        return repo_manager

sage/index.py

@@ -36,6 +36,10 @@ def main():
     for validator in arg_validators:
         validator(args)
 
+    if args.llm_retriever:
+        logging.warning("The LLM retriever does not require indexing, so this script is a no-op.")
+        return
+
     # Additionally validate embedder and vector store compatibility.
     if args.embedding_provider == "openai" and args.vector_store_provider != "pinecone":
         parser.error("When using OpenAI embedder, the vector store type must be Pinecone.")
@@ -50,22 +54,7 @@ def main():
     repo_embedder = None
     if args.index_repo:
         logging.info("Cloning the repository...")
-        repo_manager = GitHubRepoManager(
-            args.repo_id,
-            commit_hash=args.commit_hash,
-            access_token=os.getenv("GITHUB_TOKEN"),
-            local_dir=args.local_dir,
-            inclusion_file=args.include,
-            exclusion_file=args.exclude,
-        )
-
-        success = repo_manager.download()
-        if not success:
-            raise ValueError(
-                f"Unable to clone {args.repo_id}. Please check that it exists and you have access to it. "
-                "For private repositories, please set the GITHUB_TOKEN variable in your environment."
-            )
-
+        repo_manager = GitHubRepoManager.from_args(args)
         logging.info("Embedding the repo...")
         chunker = UniversalFileChunker(max_tokens=args.tokens_per_chunk)
         repo_embedder = build_batch_embedder_from_flags(repo_manager, chunker, args)

sage/retriever.py

@@ -1,32 +1,221 @@
-from typing import Optional
+import logging
+import os
+from typing import List, Optional
 
+import anthropic
+import Levenshtein
+from anytree import Node, RenderTree
+from langchain.callbacks.manager import CallbackManagerForRetrieverRun
 from langchain.retrievers import ContextualCompressionRetriever
 from langchain.retrievers.multi_query import MultiQueryRetriever
+from langchain.schema import BaseRetriever, Document
+from langchain_core.output_parsers import CommaSeparatedListOutputParser
 from langchain_google_genai import GoogleGenerativeAIEmbeddings
 from langchain_openai import OpenAIEmbeddings
 from langchain_voyageai import VoyageAIEmbeddings
+from pydantic import Field
 
-from sage.data_manager import DataManager
+from sage.data_manager import DataManager, GitHubRepoManager
 from sage.llm import build_llm_via_langchain
 from sage.reranker import build_reranker
 from sage.vector_store import build_vector_store_from_args
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+
+class LLMRetriever(BaseRetriever):
+    """Custom Langchain retriever based on an LLM.
+
+    Builds a representation of the folder structure of the repo, feeds it to an LLM, and asks the LLM for the most
+    relevant files for a particular user query, expecting it to make decisions based solely on file names.
+
+    Only works with Claude/Anthropic, because it's very slow (e.g. 15s for a mid-sized codebase) and we need prompt
+    caching to make it usable.
+    """
+
+    repo_manager: GitHubRepoManager = Field(...)
+    top_k: int = Field(...)
+    all_repo_files: List[str] = Field(...)
+    repo_hierarchy: str = Field(...)
+
+    def __init__(self, repo_manager: GitHubRepoManager, top_k: int):
+        super().__init__()
+        self.repo_manager = repo_manager
+        self.top_k = top_k
+
+        # Best practice would be to make these fields @cached_property, but that impedes class serialization.
+        self.all_repo_files = [metadata["file_path"] for metadata in self.repo_manager.walk(get_content=False)]
+        self.repo_hierarchy = LLMRetriever._render_file_hierarchy(self.all_repo_files)
+
+        if not os.environ.get("ANTHROPIC_API_KEY"):
+            raise ValueError("Please set the ANTHROPIC_API_KEY environment variable for the LLMRetriever.")
+
+    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
+        """Retrieve relevant documents for a given query."""
+        filenames = self._ask_llm_to_retrieve(user_query=query, top_k=self.top_k)
+        documents = []
+        for filename in filenames:
+            document = Document(
+                page_content=self.repo_manager.read_file(filename),
+                metadata={"file_path": filename, "url": self.repo_manager.url_for_file(filename)},
+            )
+            documents.append(document)
+        return documents
+
+    def _ask_llm_to_retrieve(self, user_query: str, top_k: int) -> List[str]:
+        """Feeds the file hierarchy and user query to the LLM and asks which files might be relevant."""
+        sys_prompt = f"""
+You are a retriever system. You will be given a user query and a list of files in a GitHub repository. Your task is to determine the top {top_k} files that are most relevant to the user query.
+DO NOT RESPOND TO THE USER QUERY DIRECTLY. Instead, respond with full paths to relevant files that could contain the answer to the query. Say absolutely nothing else other than the file paths.
+
+Here is the file hierarchy of the GitHub repository:
+
+{self.repo_hierarchy}
+"""
+        # We are deliberately repeating the "DO NOT RESPOND TO THE USER QUERY DIRECTLY" instruction here.
+        augmented_user_query = f"""
+User query: {user_query}
+
+DO NOT RESPOND TO THE USER QUERY DIRECTLY. Instead, respond with full paths to relevant files that could contain the answer to the query. Say absolutely nothing else other than the file paths.
+"""
+        response = LLMRetriever._call_via_anthropic_with_prompt_caching(sys_prompt, augmented_user_query)
+        files_from_llm = response.content[0].text.strip().split("\n")
+        validated_files = []
+
+        for filename in files_from_llm:
+            if filename not in self.all_repo_files:
+                if "/" not in filename:
+                    # This is most likely some natural language excuse from the LLM; skip it.
+                    continue
+                # Try a few heuristics to fix the filename.
+                filename = LLMRetriever._fix_filename(filename, self.repo_manager.repo_id)
+                if filename not in self.all_repo_files:
+                    # The heuristics failed; try to find the closest filename in the repo.
+                    filename = LLMRetriever._find_closest_filename(filename, self.all_repo_files)
+            if filename in self.all_repo_files:
+                validated_files.append(filename)
+        return validated_files
+
+    @staticmethod
+    def _call_via_anthropic_with_prompt_caching(system_prompt: str, user_prompt: str) -> str:
+        """Calls the Anthropic API with prompt caching for the system prompt.
+
+        See https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching.
+
+        We're circumventing LangChain for now, because the feature is < 1 week old at the time of writing and has no
+        documentation: https://github.com/langchain-ai/langchain/pull/27087
+        """
+        CLAUDE_MODEL = "claude-3-5-sonnet-20240620"
+        client = anthropic.Anthropic()
+
+        system_message = {"type": "text", "text": system_prompt, "cache_control": {"type": "ephemeral"}}
+        user_message = {"role": "user", "content": user_prompt}
+
+        response = client.beta.prompt_caching.messages.create(
+            model=CLAUDE_MODEL,
+            max_tokens=1024,  # The maximum number of *output* tokens to generate.
+            system=[system_message],
+            messages=[user_message],
+        )
+        # Caching information will be under `cache_creation_input_tokens` and `cache_read_input_tokens`.
+        # Note that, for prompts shorter than 1024 tokens, Anthropic will not do any caching.
+        logging.info("Anthropic prompt caching info: %s", response.usage)
+        return response
+
+    @staticmethod
+    def _render_file_hierarchy(file_paths: List[str]) -> str:
+        """Given a list of files, produces a visualization of the file hierarchy. For instance:
+        folder1
+            folder11
+                file111.py
+                file112.py
+            folder12
+                file121.py
+        folder2
+            file21.py
+        """
+        # The "nodepath" is the path from root to the node (e.g. huggingface/transformers/examples)
+        nodepath_to_node = {}
+
+        for path in file_paths:
+            items = path.split("/")
+            nodepath = ""
+            parent_node = None
+            for item in items:
+                nodepath = f"{nodepath}/{item}"
+                if nodepath in nodepath_to_node:
+                    node = nodepath_to_node[nodepath]
+                else:
+                    node = Node(item, parent=parent_node)
+                    nodepath_to_node[nodepath] = node
+                parent_node = node
+
+        root_path = f"/{file_paths[0].split('/')[0]}"
+        full_render = ""
+        root_node = nodepath_to_node[root_path]
+        for pre, fill, node in RenderTree(root_node):
+            render = "%s%s\n" % (pre, node.name)
+            # Replace special lines with empty strings to save on tokens.
+            render = render.replace("└", " ").replace("├", " ").replace("│", " ").replace("─", " ")
+            full_render += render
+        return full_render
+
+    @staticmethod
+    def _fix_filename(filename: str, repo_id: str) -> str:
+        """Attempts to "fix" a filename output by the LLM.
+
+        Common issues with LLM-generated filenames:
+        - The LLM prepends an extraneous "/".
+        - The LLM omits the name of the org (e.g. "transformers/README.md" instead of "huggingface/transformers/README.md").
+        - The LLM omits the name of the repo (e.g. "huggingface/README.md" instead of "huggingface/transformers/README.md").
+        - The LLM omits the org/repo prefix (e.g. "README.md" instead of "huggingface/transformers/README.md").
+        """
+        if filename.startswith("/"):
+            filename = filename[1:]
+        org_name, repo_name = repo_id.split("/")
+        items = filename.split("/")
+        if filename.startswith(org_name) and not filename.startswith(repo_id):
+            new_items = [org_name, repo_name] + items[1:]
+            return "/".join(new_items)
+        if not filename.startswith(org_name) and filename.startswith(repo_name):
+            return f"{org_name}/{filename}"
+        if not filename.startswith(org_name) and not filename.startswith(repo_name):
+            return f"{org_name}/{repo_name}/{filename}"
+        return filename
+
+    @staticmethod
+    def _find_closest_filename(filename: str, repo_filenames: List[str], max_edit_distance: int = 10) -> Optional[str]:
+        """Returns the path in the repo with smallest edit distance from `filename`. Helpful when the `filename` was
+        generated by an LLM and parts of it might have been hallucinated. Returns None if the closest path is more than
+        `max_edit_distance` away. In case of a tie, returns an arbitrary closest path.
+        """
+        distances = [(path, Levenshtein.distance(filename, path)) for path in repo_filenames]
+        distances.sort(key=lambda x: x[1])
+        if distances[0][1] <= max_edit_distance:
+            closest_path = distances[0][0]
+            return closest_path
+        return None
+
 
 def build_retriever_from_args(args, data_manager: Optional[DataManager] = None):
     """Builds a retriever (with optional reranking) from command-line arguments."""
-
-    if args.embedding_provider == "openai":
-        embeddings = OpenAIEmbeddings(model=args.embedding_model)
-    elif args.embedding_provider == "voyage":
-        embeddings = VoyageAIEmbeddings(model=args.embedding_model)
-    elif args.embedding_provider == "gemini":
-        embeddings = GoogleGenerativeAIEmbeddings(model=args.embedding_model)
+    if args.llm_retriever:
+        retriever = LLMRetriever(GitHubRepoManager.from_args(args), top_k=args.retriever_top_k)
     else:
-        embeddings = None
+        if args.embedding_provider == "openai":
+            embeddings = OpenAIEmbeddings(model=args.embedding_model)
+        elif args.embedding_provider == "voyage":
+            embeddings = VoyageAIEmbeddings(model=args.embedding_model)
+        elif args.embedding_provider == "gemini":
+            embeddings = GoogleGenerativeAIEmbeddings(model=args.embedding_model)
+        else:
+            embeddings = None
 
-    retriever = build_vector_store_from_args(args, data_manager).as_retriever(
-        top_k=args.retriever_top_k, embeddings=embeddings, namespace=args.index_namespace
-    )
+        retriever = build_vector_store_from_args(args, data_manager).as_retriever(
+            top_k=args.retriever_top_k, embeddings=embeddings, namespace=args.index_namespace
+        )
 
     if args.multi_query_retriever:
         retriever = MultiQueryRetriever.from_llm(