fix merge conflict

README.md

@@ -141,6 +141,15 @@ If you are planning on indexing GitHub issues in addition to the codebase, you w
 
 ## Additional features
 
+<details>
+<summary><strong>:lock: Working with private repositories</strong></summary>
+To index and chat with a private repository, simply set the GITHUB_TOKEN environment variable. To obtain this token: go to github.com > click on your profile icon > Settings > Developer settings > Personal access tokens. You can either make a fine-grained token for the desired repository, or a classic token.
+
+```
+export GITHUB_TOKEN=...
+```
+</details>
+
 <details>
 <summary><strong>:hammer_and_wrench: Control which files get indexed</strong></summary>
 
@@ -177,6 +186,7 @@ By default, we use the exclusion file [sample-exclude.txt](sage/sample-exclude.t
 <details>
 <summary><strong>:bug: Index open GitHub issues</strong></summary>
 You will need a GitHub token first:
+
 ```
 export GITHUB_TOKEN=...
 ```

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,11 +1,10 @@
 """Runs a batch job to compute embeddings for an entire repo and stores them into a vector store."""
 
 import logging
-import nltk
+import os
 import time
 
 import configargparse
-import pkg_resources
 
 import sage.config as sage_config
 from sage.chunker import UniversalFileChunker
@@ -55,6 +54,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,
@@ -69,7 +69,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)
@@ -94,7 +96,7 @@ def main():
         logging.info("Moving embeddings to the repo vector store...")
         repo_vector_store = build_vector_store_from_args(args, repo_manager)
         repo_vector_store.ensure_exists()
-        repo_vector_store.upsert(repo_embedder.download_embeddings(repo_jobs_file))
+        repo_vector_store.upsert(repo_embedder.download_embeddings(repo_jobs_file), namespace=args.index_namespace)
 
     if issues_embedder is not None:
         logging.info("Waiting for issue embeddings to be ready...")
@@ -105,7 +107,9 @@ def main():
         logging.info("Moving embeddings to the issues vector store...")
         issues_vector_store = build_vector_store_from_args(args, issues_manager)
         issues_vector_store.ensure_exists()
-        issues_vector_store.upsert(issues_embedder.download_embeddings(issues_jobs_file))
+        issues_vector_store.upsert(
+            issues_embedder.download_embeddings(issues_jobs_file), namespace=args.index_namespace
+        )
 
     logging.info("Done!")
 

sage/retriever.py

@@ -2,7 +2,6 @@ from langchain.retrievers import ContextualCompressionRetriever
 from langchain_openai import OpenAIEmbeddings
 from langchain_voyageai import VoyageAIEmbeddings
 
-
 from sage.reranker import build_reranker
 from sage.vector_store import build_vector_store_from_args
 
@@ -17,7 +16,9 @@ def build_retriever_from_args(args):
     else:
         embeddings = None
 
-    retriever = build_vector_store_from_args(args).as_retriever(top_k=args.retriever_top_k, embeddings=embeddings)
+    retriever = build_vector_store_from_args(args).as_retriever(
+        top_k=args.retriever_top_k, embeddings=embeddings, namespace=args.index_namespace
+    )
 
     reranker = build_reranker(args.reranker_provider, args.reranker_model, args.reranker_top_k)
     if reranker:

sage/vector_store.py

@@ -1,7 +1,7 @@
 """Vector store abstraction and implementations."""
 
-import os
 import logging
+import os
 from abc import ABC, abstractmethod
 from functools import cached_property
 from typing import Dict, Generator, List, Optional, Tuple
@@ -36,33 +36,32 @@ class VectorStore(ABC):
         """Ensures that the vector store exists. Creates it if it doesn't."""
 
     @abstractmethod
-    def upsert_batch(self, vectors: List[Vector]):
+    def upsert_batch(self, vectors: List[Vector], namespace: str):
         """Upserts a batch of vectors."""
 
-    def upsert(self, vectors: Generator[Vector, None, None]):
+    def upsert(self, vectors: Generator[Vector, None, None], namespace: str):
         """Upserts in batches of 100, since vector stores have a limit on upsert size."""
         batch = []
         for metadata, embedding in vectors:
             batch.append((metadata, embedding))
             if len(batch) == 100:
-                self.upsert_batch(batch)
+                self.upsert_batch(batch, namespace)
                 batch = []
         if batch:
-            self.upsert_batch(batch)
+            self.upsert_batch(batch, namespace)
 
     @abstractmethod
-    def as_retriever(self, top_k: int, embeddings: Embeddings):
+    def as_retriever(self, top_k: int, embeddings: Embeddings, namespace: str):
         """Converts the vector store to a LangChain retriever object."""
 
 
 class PineconeVectorStore(VectorStore):
     """Vector store implementation using Pinecone."""
 
-    def __init__(self, index_name: str, namespace: str, dimension: int, alpha: float, bm25_cache: Optional[str] = None):
+    def __init__(self, index_name: str, dimension: int, alpha: float, bm25_cache: Optional[str] = None):
         """
         Args:
             index_name: The name of the Pinecone index to use. If it doesn't exist already, we'll create it.
-            namespace: The namespace within the index to use.
             dimension: The dimension of the vectors.
             alpha: The alpha parameter for hybrid search: alpha == 1.0 means pure dense search, alpha == 0.0 means pure
                 BM25, and 0.0 < alpha < 1.0 means a hybrid of the two.
@@ -72,7 +71,6 @@ class PineconeVectorStore(VectorStore):
         self.index_name = index_name
         self.dimension = dimension
         self.client = Pinecone()
-        self.namespace = namespace
         self.alpha = alpha
 
         if alpha < 1.0:
@@ -105,7 +103,8 @@ class PineconeVectorStore(VectorStore):
         def patched_query(*args, **kwargs):
             result = original_query(*args, **kwargs)
             for res in result["matches"]:
-                res["metadata"]["context"] = res["metadata"][TEXT_FIELD]
+                if TEXT_FIELD in res["metadata"]:
+                    res["metadata"]["context"] = res["metadata"][TEXT_FIELD]
             return result
 
         index.query = patched_query
@@ -121,7 +120,7 @@ class PineconeVectorStore(VectorStore):
                 spec=ServerlessSpec(cloud="aws", region="us-east-1"),
             )
 
-    def upsert_batch(self, vectors: List[Vector]):
+    def upsert_batch(self, vectors: List[Vector], namespace: str):
         pinecone_vectors = []
         for i, (metadata, embedding) in enumerate(vectors):
             vector = {"id": metadata.get("id", str(i)), "values": embedding, "metadata": metadata}
@@ -129,21 +128,21 @@ class PineconeVectorStore(VectorStore):
                 vector["sparse_values"] = self.bm25_encoder.encode_documents(metadata[TEXT_FIELD])
             pinecone_vectors.append(vector)
 
-        self.index.upsert(vectors=pinecone_vectors, namespace=self.namespace)
+        self.index.upsert(vectors=pinecone_vectors, namespace=namespace)
 
-    def as_retriever(self, top_k: int, embeddings: Embeddings):
+    def as_retriever(self, top_k: int, embeddings: Embeddings, namespace: str):
         if self.bm25_encoder:
             return PineconeHybridSearchRetriever(
                 embeddings=embeddings,
                 sparse_encoder=self.bm25_encoder,
                 index=self.index,
-                namespace=self.namespace,
+                namespace=namespace,
                 top_k=top_k,
                 alpha=self.alpha,
             )
 
         return LangChainPinecone.from_existing_index(
-            index_name=self.index_name, embedding=embeddings, namespace=self.namespace
+            index_name=self.index_name, embedding=embeddings, namespace=namespace
         ).as_retriever(search_kwargs={"k": top_k})
 
 
@@ -157,12 +156,14 @@ class MarqoVectorStore(VectorStore):
     def ensure_exists(self):
         pass
 
-    def upsert_batch(self, vectors: List[Vector]):
+    def upsert_batch(self, vectors: List[Vector], namespace: str):
         # Since Marqo is both an embedder and a vector store, the embedder is already doing the upsert.
         pass
 
-    def as_retriever(self, top_k: int, embeddings: Embeddings = None):
+    def as_retriever(self, top_k: int, embeddings: Embeddings = None, namespace: str = None):
         del embeddings  # Unused; The Marqo vector store is also an embedder.
+        del namespace  # Unused; Unlike Pinecone, Marqo doesn't differentiate between index name and namespace.
+
         vectorstore = Marqo(client=self.client, index_name=self.index_name)
 
         # Monkey-patch the _construct_documents_from_results_without_score method to not expect a "metadata" field in
@@ -202,7 +203,6 @@ def build_vector_store_from_args(args: dict, data_manager: Optional[DataManager]
 
         return PineconeVectorStore(
             index_name=args.pinecone_index_name,
-            namespace=args.index_namespace,
             dimension=args.embedding_size if "embedding_size" in args else None,
             alpha=args.retrieval_alpha,
             bm25_cache=bm25_cache,

setup.py

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