Take "namespace" out of the vector store constructors.

sage/index.py

@@ -92,7 +92,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...")
@@ -103,7 +103,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
@@ -29,33 +29,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.
@@ -65,7 +64,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:
@@ -107,7 +105,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}
@@ -115,21 +113,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})
 
 
@@ -143,12 +141,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
@@ -188,7 +188,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,