Pass access_token in the constructor of DataManagers.

sage/data_manager.py

@@ -30,6 +30,7 @@ class GitHubRepoManager(DataManager):
         self,
         repo_id: str,
         commit_hash: str = None,
+        access_token: str = None,
         local_dir: str = None,
         inclusion_file: str = None,
         exclusion_file: str = None,
@@ -38,6 +39,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.
+            access_token: A GitHub access token to use for cloning private repositories. Not needed for public repos.
             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".
@@ -47,6 +49,7 @@ class GitHubRepoManager(DataManager):
         super().__init__(dataset_id=repo_id)
         self.repo_id = repo_id
         self.commit_hash = commit_hash
+        self.access_token = access_token
 
         self.local_dir = local_dir or "/tmp/"
         if not os.path.exists(self.local_dir):
@@ -57,8 +60,6 @@ class GitHubRepoManager(DataManager):
         if not os.path.exists(self.log_dir):
             os.makedirs(self.log_dir)
 
-        self.access_token = os.getenv("GITHUB_TOKEN")
-
         if inclusion_file and exclusion_file:
             raise ValueError("Only one of inclusion_file or exclusion_file should be provided.")
 

sage/github.py

@@ -1,7 +1,6 @@
 """GitHub-specific implementations for DataManager and Chunker."""
 
 import logging
-import os
 from dataclasses import dataclass
 from typing import Any, Dict, Generator, List, Tuple
 
@@ -47,12 +46,12 @@ class GitHubIssue:
 class GitHubIssuesManager(DataManager):
     """Class to manage the GitHub issues of a particular repository."""
 
-    def __init__(self, repo_id: str, index_comments: bool = False, max_issues: int = None):
+    def __init__(self, repo_id: str, access_token: str, index_comments: bool = False, max_issues: int = None):
         super().__init__(dataset_id=repo_id + "/issues")
         self.repo_id = repo_id
         self.index_comments = index_comments
         self.max_issues = max_issues
-        self.access_token = os.getenv("GITHUB_TOKEN")
+        self.access_token = access_token
         if not self.access_token:
             raise ValueError("Please set the GITHUB_TOKEN environment variable when indexing GitHub issues.")
         self.issues = []

sage/index.py

@@ -1,10 +1,10 @@
 """Runs a batch job to compute embeddings for an entire repo and stores them into a vector store."""
 
 import logging
+import os
 import time
 
 import configargparse
-import pkg_resources
 
 import sage.config as sage_config
 from sage.chunker import UniversalFileChunker
@@ -53,6 +53,7 @@ def main():
         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,
@@ -67,7 +68,9 @@ def main():
     issues_embedder = None
     if args.index_issues:
         logging.info("Issuing embedding jobs for GitHub issues...")
-        issues_manager = GitHubIssuesManager(args.repo_id, index_comments=args.index_issue_comments)
+        issues_manager = GitHubIssuesManager(
+            args.repo_id, access_token=os.getenv("GITHUB_TOKEN"), index_comments=args.index_issue_comments
+        )
         issues_manager.download()
         logging.info("Embedding GitHub issues...")
         chunker = GitHubIssuesChunker(max_tokens=args.tokens_per_chunk)

setup.py

@@ -8,7 +8,7 @@ def readfile(filename):
 
 setup(
     name="sage",
-    version="0.1.1",
+    version="0.1.2",
     packages=find_packages(),
     include_package_data=True,
     package_data={