Option to index repo at specific commit.

sage/data_manager.py

@@ -29,6 +29,7 @@ class GitHubRepoManager(DataManager):
     def __init__(
         self,
         repo_id: str,
+        commit_hash: str = None,
         local_dir: str = None,
         inclusion_file: str = None,
         exclusion_file: str = None,
@@ -36,6 +37,7 @@ class GitHubRepoManager(DataManager):
         """
         Args:
             repo_id: The identifier of the repository in owner/repo format, e.g. "Storia-AI/sage".
+            commit_hash: Optional commit hash to checkout. If not specified, we pull the latest version of the repo.
             local_dir: The local directory where the repository will be cloned.
             inclusion_file: A file with a lists of files/directories/extensions to include. Each line must be in one of
                 the following formats: "ext:.my-extension", "file:my-file.py", or "dir:my-directory".
@@ -44,6 +46,7 @@ class GitHubRepoManager(DataManager):
         """
         super().__init__(dataset_id=repo_id)
         self.repo_id = repo_id
+        self.commit_hash = commit_hash
 
         self.local_dir = local_dir or "/tmp/"
         if not os.path.exists(self.local_dir):
@@ -103,7 +106,11 @@ class GitHubRepoManager(DataManager):
             clone_url = f"https://github.com/{self.repo_id}.git"
 
         try:
-            Repo.clone_from(clone_url, self.local_path, depth=1, single_branch=True)
+            if self.commit_hash:
+                repo = Repo.clone_from(clone_url, self.local_path)
+                repo.git.checkout(self.commit_hash)
+            else:
+                Repo.clone_from(clone_url, self.local_path, depth=1, single_branch=True)
         except GitCommandError as e:
             logging.error("Unable to clone %s from %s. Error: %s", self.repo_id, clone_url, e)
             return False

sage/index.py

@@ -38,6 +38,7 @@ OPENAI_DEFAULT_EMBEDDING_SIZE = {
 def main():
     parser = argparse.ArgumentParser(description="Batch-embeds a GitHub repository and its issues.")
     parser.add_argument("repo_id", help="The ID of the repository to index")
+    parser.add_argument("--commit-hash", help="Optional commit hash to checkout. When not provided, defaults to HEAD.")
     parser.add_argument("--embedder-type", default="marqo", choices=["openai", "marqo"])
     parser.add_argument(
         "--embedding-model",
@@ -72,9 +73,8 @@ def main():
     parser.add_argument(
         "--index-name",
         default=None,
-        help="Vector store index name. For Marqo, we default it to the repository name. Required for Pinecone, since "
-        "it needs to be created manually on their website. In Pinecone terminology, this is *not* the namespace (which "
-        "we default to the repo ID).",
+        help="Vector store index name. For Marqo, we default it to the repository name. Required for Pinecone. "
+        "In Pinecone terminology, this is *not* the namespace (which we default to the repo ID).",
     )
     parser.add_argument(
         "--include",
@@ -202,6 +202,7 @@ def main():
         logging.info("Cloning the repository...")
         repo_manager = GitHubRepoManager(
             args.repo_id,
+            commit_hash=args.commit_hash,
             local_dir=args.local_dir,
             inclusion_file=args.include,
             exclusion_file=args.exclude,

sage/vector_store.py

@@ -152,12 +152,24 @@ def build_from_args(args: dict) -> VectorStore:
         if not args.index_name:
             raise ValueError("Please specify --index-name for Pinecone.")
         dimension = args.embedding_size if "embedding_size" in args else None
+
+        namespace = args.repo_id
+        if args.commit_hash:
+            namespace += "/" + args.commit_hash
+
         return PineconeVectorStore(
-            index_name=args.index_name, namespace=args.repo_id, dimension=dimension, hybrid=args.hybrid_retrieval
+            index_name=args.index_name, namespace=namespace, dimension=dimension, hybrid=args.hybrid_retrieval
         )
     elif args.vector_store_type == "marqo":
         marqo_url = args.marqo_url or "http://localhost:8882"
-        index_name = args.index_name or args.repo_id.split("/")[1]
+
+        index_name = args.index_name
+        if not index_name:
+            # Marqo doesn't allow slashes in the index name.
+            index_name = args.repo_id.split("/")[1]
+            if args.commit_hash:
+                index_name += "_" + args.commit_hash
+
         return MarqoVectorStore(url=marqo_url, index_name=index_name)
     else:
         raise ValueError(f"Unrecognized vector store type {args.vector_store_type}")