Add YAML configurations (#38)

MANIFEST.in

@@ -1 +1,3 @@
-include sage/sample-exclude.txt
+include sage/sample-exclude.txt
+include sage/configs/local.yaml
+include sage/configs/remote.yaml

README.md

@@ -49,7 +49,7 @@ pip install git+https://github.com/Storia-AI/sage.git@main
 2. Enables chatting via LLM + RAG (requiring access to an LLM)
 
 <details open>
-<summary><strong>:computer: Running locally</strong></summary>
+<summary><strong>:computer: Running locally (lower quality)</strong></summary>
 
 1. To index the codebase locally, we use the open-source project <a href="https://github.com/marqo-ai/marqo">Marqo</a>, which is both an embedder and a vector store. To bring up a Marqo instance:
 
@@ -70,7 +70,7 @@ pip install git+https://github.com/Storia-AI/sage.git@main
 </details>
 
 <details>
-<summary><strong>:cloud: Using external providers</strong></summary>
+<summary><strong>:cloud: Using external providers (higher quality)</strong></summary>
 
 1. We support <a href="https://openai.com/">OpenAI</a> for embeddings (they have a super fast batch embedding API) and <a href="https://www.pinecone.io/">Pinecone</a> for the vector store. So you will need two API keys:
 
@@ -84,37 +84,27 @@ pip install git+https://github.com/Storia-AI/sage.git@main
     export PINECONE_INDEX_NAME=...
     ```
 
-2. For chatting with an LLM, we support OpenAI and Anthropic. For the latter, set an additional API key:
-
+3. For reranking, we use <a href="https://cohere.com/rerank">Cohere</a> by default, but you can also try rerankers from <a href="https://developer.nvidia.com/blog/enhancing-rag-pipelines-with-re-ranking/">NVIDIA</a> or <a href="https://jina.ai/reranker/">Jina</a>:
     ```
-    export ANTHROPIC_API_KEY=...
+    export COHERE_API_KEY=...  # or
+    export NVIDIA_API_KEY=...  # or
+    export JINA_API_KEY=...
     ```
 
-</details>
-
-<br>
-<summary><strong>Optional</strong></summary>
-
-- By default, we use an <a href="https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-6-v2">open-source re-ranker</a>. For higher accuracy, you can use <a href="https://cohere.com/rerank">Cohere</a>, <a href="https://developer.nvidia.com/blog/enhancing-rag-pipelines-with-re-ranking/">NVIDIA</a> or <a href="https://jina.ai/reranker/">Jina</a>:
-
+4. For chatting with an LLM, we support OpenAI and Anthropic. For the latter, set an additional API key:
     ```
-    export COHERE_API_KEY=...
-    export NVIDIA_API_KEY=...
-    export JINA_API_KEY=...
+    export ANTHROPIC_API_KEY=...
     ```
 
-    We are seeing significant gains in accuracy from these proprietary rerankers.
+</details>
 
-- If you are planning on indexing GitHub issues in addition to the codebase, you will need a GitHub token:
+### Optional
+If you are planning on indexing GitHub issues in addition to the codebase, you will need a GitHub token:
 
     export GITHUB_TOKEN=...
 
-
 ## Running it
 
-<details open>
-<summary><strong>:computer: Run locally</strong></summary>
-
 1. Select your desired repository:
     ```
     export GITHUB_REPO=huggingface/transformers
@@ -124,41 +114,18 @@ pip install git+https://github.com/Storia-AI/sage.git@main
     ```
     sage-index $GITHUB_REPO
     ```
+    To use external providers instead of running locally, set `--mode=remote`.
 
 3. Chat with the repository, once it's indexed:
     ```
     sage-chat $GITHUB_REPO
     ```
-    To get a public URL for your chat app, set `--share=true`.
-
+    To use external providers instead of running locally, set `--mode=remote`.
 </details>
 
-<details>
-<summary><strong>:cloud: Use external providers</strong></summary>
-
-1. Select your desired repository:
-    ```
-    export GITHUB_REPO=huggingface/transformers
-    ```
-
-2. Index the repository. This might take a few minutes, depending on its size.
-    ```
-    sage-index $GITHUB_REPO \
-        --embedder-type=openai \
-        --vector-store=pinecone \
-        --index-name=$PINECONE_INDEX_NAME
-    ```
-
-3. Chat with the repository, once it's indexed:
-    ```
-    sage-chat $GITHUB_REPO \
-        --vector-store-type=pinecone \
-        --index-name=$PINECONE_INDEX_NAME \
-        --llm-provider=openai \
-        --llm-model=gpt-4
-    ```
-    To get a public URL for your chat app, set `--share=true`.
-</details>
+### Notes:
+- To get a public URL for your chat app, set `--share=true`.
+- You can overwrite the default settings (e.g. desired embedding model or LLM) via command line flags. Run `sage-index --help` or `sage-chat --help` for a full list.
 
 ## Additional features
 

sage/chat.py

@@ -3,11 +3,11 @@
 You must run `sage-index $GITHUB_REPO` first in order to index the codebase into a vector store.
 """
 
-import argparse
 import logging
-import os
 
+import configargparse
 import gradio as gr
+import pkg_resources
 from dotenv import load_dotenv
 from langchain.chains import create_history_aware_retriever, create_retrieval_chain
 from langchain.chains.combine_documents import create_stuff_documents_chain
@@ -15,9 +15,10 @@ from langchain.retrievers import ContextualCompressionRetriever
 from langchain.schema import AIMessage, HumanMessage
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 
-import sage.vector_store as vector_store
+import sage.config as sage_config
 from sage.llm import build_llm_via_langchain
-from sage.reranker import build_reranker, RerankerProvider
+from sage.reranker import build_reranker
+from sage.vector_store import build_vector_store_from_args
 
 load_dotenv()
 
@@ -27,7 +28,7 @@ def build_rag_chain(args):
     llm = build_llm_via_langchain(args.llm_provider, args.llm_model)
 
     retriever_top_k = 5 if args.reranker_provider == "none" else 25
-    retriever = vector_store.build_from_args(args).as_retriever(top_k=retriever_top_k)
+    retriever = build_vector_store_from_args(args).as_retriever(top_k=retriever_top_k)
     compressor = build_reranker(args.reranker_provider, args.reranker_model)
     if compressor:
         retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=retriever)
@@ -70,49 +71,27 @@ def build_rag_chain(args):
 
 
 def main():
-    parser = argparse.ArgumentParser(description="UI to chat with your codebase")
-    parser.add_argument("repo_id", help="The ID of the repository to index")
-    parser.add_argument("--llm-provider", default="ollama", choices=["openai", "anthropic", "ollama"])
-    parser.add_argument(
-        "--llm-model",
-        help="The LLM name. Must be supported by the provider specified via --llm-provider.",
+    parser = configargparse.ArgParser(
+        description="Batch-embeds a GitHub repository and its issues.", ignore_unknown_config_file_keys=True
     )
-    parser.add_argument("--vector-store-type", default="marqo", choices=["pinecone", "marqo"])
-    parser.add_argument("--index-name", help="Vector store index name. Required for Pinecone.")
-    parser.add_argument(
-        "--marqo-url",
-        default="http://localhost:8882",
-        help="URL for the Marqo server. Required if using Marqo as embedder or vector store.",
-    )
-    parser.add_argument("--reranker-provider", default="huggingface", choices=[r.value for r in RerankerProvider])
-    parser.add_argument(
-        "--reranker-model",
-        help="The reranker model name. When --reranker-provider=huggingface, we suggest choosing a model from the "
-        "SentenceTransformers Cross-Encoders library https://huggingface.co/cross-encoder?sort_models=downloads#models",
-    )
-    parser.add_argument(
+    parser.add(
         "--share",
         default=False,
         help="Whether to make the gradio app publicly accessible.",
     )
-    parser.add_argument(
-        "--hybrid-retrieval",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Whether to use a hybrid of vector DB + BM25 retrieval. When set to False, we only use vector DB "
-        "retrieval. This is only relevant if using Pinecone as the vector store.",
-    )
+    sage_config.add_config_args(parser)
+
+    arg_validators = [
+        sage_config.add_repo_args(parser),
+        sage_config.add_vector_store_args(parser),
+        sage_config.add_reranking_args(parser),
+        sage_config.add_llm_args(parser),
+    ]
+
     args = parser.parse_args()
 
-    if not args.llm_model:
-        if args.llm_provider == "openai":
-            args.llm_model = "gpt-4"
-        elif args.llm_provider == "anthropic":
-            args.llm_model = "claude-3-opus-20240229"
-        elif args.llm_provider == "ollama":
-            args.llm_model = "llama3.1"
-        else:
-            raise ValueError("Please specify --llm_model")
+    for validator in arg_validators:
+        validator(args)
 
     rag_chain = build_rag_chain(args)
 

sage/config.py

@@ -0,0 +1,293 @@
+"""Utility methods to define and validate flags."""
+
+import argparse
+import logging
+import os
+import re
+from typing import Callable
+
+import pkg_resources
+from configargparse import ArgumentParser
+
+from sage.reranker import RerankerProvider
+
+MARQO_MAX_CHUNKS_PER_BATCH = 64
+# The ADA embedder from OpenAI has a maximum of 8192 tokens.
+OPENAI_MAX_TOKENS_PER_CHUNK = 8192
+# The OpenAI batch embedding API enforces a maximum of 2048 chunks per batch.
+OPENAI_MAX_CHUNKS_PER_BATCH = 2048
+# The OpenAI batch embedding API enforces a maximum of 3M tokens processed at once.
+OPENAI_MAX_TOKENS_PER_JOB = 3_000_000
+
+# Note that OpenAI embedding models have fixed dimensions, however, taking a slice of them is possible.
+# See "Reducing embedding dimensions" under https://platform.openai.com/docs/guides/embeddings/use-cases and
+# https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-dimensions
+OPENAI_DEFAULT_EMBEDDING_SIZE = {
+    "text-embedding-ada-002": 1536,
+    "text-embedding-3-small": 1536,
+    "text-embedding-3-large": 3072,
+}
+
+
+def add_config_args(parser: ArgumentParser):
+    """Adds configuration-related arguments to the parser."""
+    parser.add(
+        "--mode",
+        choices=["local", "remote"],
+        default="local",
+        help="Whether to use local-only resources or call third-party providers.",
+    )
+    parser.add(
+        "--config",
+        is_config_file=True,
+        help="Path to .yaml configuration file.",
+    )
+    args, _ = parser.parse_known_args()
+    config_file = pkg_resources.resource_filename(__name__, f"configs/{args.mode}.yaml")
+    parser.set_defaults(config=config_file)
+
+
+def add_repo_args(parser: ArgumentParser) -> Callable:
+    """Adds repository-related arguments to the parser and returns a validator."""
+    parser.add("repo_id", help="The ID of the repository to index")
+    parser.add("--commit-hash", help="Optional commit hash to checkout. When not provided, defaults to HEAD.")
+    parser.add(
+        "--local-dir",
+        default="repos",
+        help="The local directory to store the repository",
+    )
+    return validate_repo_args
+
+
+def add_embedding_args(parser: ArgumentParser) -> Callable:
+    """Adds embedding-related arguments to the parser and returns a validator."""
+    parser.add("--embedding-provider", default="marqo", choices=["openai", "marqo"])
+    parser.add(
+        "--embedding-model",
+        type=str,
+        default=None,
+        help="The embedding model. Defaults to `text-embedding-ada-002` for OpenAI and `hf/e5-base-v2` for Marqo.",
+    )
+    parser.add(
+        "--embedding-size",
+        type=int,
+        default=None,
+        help="The embedding size to use for OpenAI text-embedding-3* models. Defaults to 1536 for small and 3072 for "
+        "large. Note that no other OpenAI models support a dynamic embedding size, nor do models used with Marqo.",
+    )
+    parser.add(
+        "--tokens-per-chunk",
+        type=int,
+        default=800,
+        help="https://arxiv.org/pdf/2406.14497 recommends a value between 200-800.",
+    )
+    parser.add(
+        "--chunks-per-batch",
+        type=int,
+        help="Maximum chunks per batch. We recommend 2000 for the OpenAI embedder. Marqo enforces a limit of 64.",
+    )
+    parser.add(
+        "--max-embedding-jobs",
+        type=int,
+        help="Maximum number of embedding jobs to run. Specifying this might result in "
+        "indexing only part of the repository, but prevents you from burning through OpenAI credits.",
+    )
+    return validate_embedding_args
+
+
+def add_vector_store_args(parser: ArgumentParser) -> Callable:
+    """Adds vector store-related arguments to the parser and returns a validator."""
+    parser.add("--vector-store-provider", default="marqo", choices=["pinecone", "marqo"])
+    parser.add(
+        "--pinecone-index-name",
+        default=None,
+        help="Pinecone index name. Required if using Pinecone as the vector store. If the index doesn't exist already, "
+        "we will create it.",
+    )
+    parser.add(
+        "--index-namespace",
+        default=None,
+        help="Index namespace for this repo. When not specified, we default it to a derivative of the repo name.",
+    )
+    parser.add(
+        "--marqo-url",
+        default="http://localhost:8882",
+        help="URL for the Marqo server. Required if using Marqo as embedder or vector store.",
+    )
+    parser.add(
+        "--hybrid-retrieval",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Whether to use a hybrid of vector DB + BM25 retrieval. When set to False, we only use vector DB "
+        "retrieval. This is only relevant if using Pinecone as the vector store.",
+    )
+    return validate_vector_store_args
+
+
+def add_indexing_args(parser: ArgumentParser) -> Callable:
+    """Adds indexing-related arguments to the parser and returns a validator."""
+    parser.add(
+        "--include",
+        help="Path to a file containing a list of extensions to include. One extension per line.",
+    )
+    parser.add(
+        "--exclude",
+        help="Path to a file containing a list of extensions to exclude. One extension per line.",
+    )
+    # Pass --no-index-repo in order to not index the repository.
+    parser.add(
+        "--index-repo",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Whether to index the repository. At least one of --index-repo and --index-issues must be True.",
+    )
+    # Pass --no-index-issues in order to not index the issues.
+    parser.add(
+        "--index-issues",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Whether to index GitHub issues. At least one of --index-repo and --index-issues must be True. When "
+        "--index-issues is set, you must also set a GITHUB_TOKEN environment variable.",
+    )
+    # Pass --no-index-issue-comments in order to not index the comments of GitHub issues.
+    parser.add(
+        "--index-issue-comments",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Whether to index the comments of GitHub issues. This is only relevant if --index-issues is set. "
+        "GitHub's API for downloading comments is quite slow. Indexing solely the body of an issue seems to bring most "
+        "of the gains anyway.",
+    )
+    return validate_indexing_args
+
+
+def add_reranking_args(parser: ArgumentParser) -> Callable:
+    """Adds reranking-related arguments to the parser."""
+    parser.add("--reranker-provider", default="huggingface", choices=[r.value for r in RerankerProvider])
+    parser.add(
+        "--reranker-model",
+        help="The reranker model name. When --reranker-provider=huggingface, we suggest choosing a model from the "
+        "SentenceTransformers Cross-Encoders library https://huggingface.co/cross-encoder?sort_models=downloads#models",
+    )
+    # Trivial validator (nothing to check).
+    return lambda _: True
+
+
+def add_llm_args(parser: ArgumentParser) -> Callable:
+    """Adds language model-related arguments to the parser."""
+    parser.add("--llm-provider", default="ollama", choices=["openai", "anthropic", "ollama"])
+    parser.add(
+        "--llm-model",
+        help="The LLM name. Must be supported by the provider specified via --llm-provider.",
+    )
+    # Trivial validator (nothing to check).
+    return lambda _: True
+
+
+def validate_repo_args(args):
+    """Validates the configuration of the repository."""
+    if not re.match(r"^[^/]+/[^/]+$", args.repo_id):
+        raise ValueError("repo_id must be in the format 'owner/repo'")
+
+
+def _validate_openai_embedding_args(args):
+    """Validates the configuration of the OpenAI batch embedder and sets defaults."""
+    if args.embedding_provider == "openai" and not os.getenv("OPENAI_API_KEY"):
+        raise ValueError("Please set the OPENAI_API_KEY environment variable.")
+
+    if not args.embedding_model:
+        args.embedding_model = "text-embedding-ada-002"
+
+    if args.embedding_model not in OPENAI_DEFAULT_EMBEDDING_SIZE.keys():
+        raise ValueError(f"Unrecognized embeddings.model={args.embedding_model}")
+
+    if not args.embedding_size:
+        args.embedding_size = OPENAI_DEFAULT_EMBEDDING_SIZE.get(args.embedding_model)
+
+    if not args.tokens_per_chunk:
+        # https://arxiv.org/pdf/2406.14497 recommends a value between 200-800.
+        args.tokens_per_chunk = 800
+    elif args.tokens_per_chunk > OPENAI_MAX_TOKENS_PER_CHUNK:
+        args.tokens_per_chunk = OPENAI_MAX_TOKENS_PER_CHUNK
+        logging.warning(
+            f"OpenAI enforces a limit of {OPENAI_MAX_TOKENS_PER_CHUNK} tokens per chunk. "
+            "Overwriting embeddings.tokens_per_chunk."
+        )
+
+    if not args.chunks_per_batch:
+        args.chunks_per_batch = OPENAI_MAX_CHUNKS_PER_BATCH
+    elif args.chunks_per_batch > OPENAI_MAX_CHUNKS_PER_BATCH:
+        args.chunks_per_batch = OPENAI_MAX_CHUNKS_PER_BATCH
+        logging.warning(
+            f"OpenAI enforces a limit of {OPENAI_MAX_CHUNKS_PER_BATCH} chunks per batch. "
+            "Overwriting embeddings.chunks_per_batch."
+        )
+
+    chunks_per_job = args.tokens_per_chunk * args.chunks_per_batch
+    if chunks_per_job >= OPENAI_MAX_TOKENS_PER_JOB:
+        raise ValueError(f"The maximum number of chunks per job is {OPENAI_MAX_TOKENS_PER_JOB}. Got {chunks_per_job}")
+
+
+def _validate_marqo_embedding_args(args):
+    """Validates the configuration of the Marqo batch embedder and sets defaults."""
+    if not args.embedding_model:
+        args.embedding_model = "hf/e5-base-v2"
+
+    if not args.chunks_per_batch:
+        args.chunks_per_batch = MARQO_MAX_CHUNKS_PER_BATCH
+    elif args.chunks_per_batch > MARQO_MAX_CHUNKS_PER_BATCH:
+        args.chunks_per_batch = MARQO_MAX_CHUNKS_PER_BATCH
+        logging.warning(
+            f"Marqo enforces a limit of {MARQO_MAX_CHUNKS_PER_BATCH} chunks per batch. "
+            "Overwriting embeddings.chunks_per_batch."
+        )
+
+
+def validate_embedding_args(args):
+    """Validates the configuration of the batch embedder and sets defaults."""
+    if args.embedding_provider == "openai":
+        _validate_openai_embedding_args(args)
+    elif args.embedding_provider == "marqo":
+        _validate_marqo_embedding_args(args)
+    else:
+        raise ValueError(f"Unrecognized --embedding-provider={args.embedding_provider}")
+
+
+def validate_vector_store_args(args):
+    """Validates the configuration of the vector store and sets defaults."""
+
+    if not args.index_namespace:
+        args.index_namespace = args.repo_id
+        if args.commit_hash:
+            args.index_namespace += "/" + args.commit_hash
+        if args.vector_store_provider == "marqo":
+            # Marqo doesn't allow slashes in the index namespace.
+            args.index_namespace = args.index_namespace.replace("/", "_")
+
+    if args.vector_store_provider == "marqo":
+        if not args.marqo_url:
+            args.marqo_url = "http://localhost:8882"
+        if "/" in args.index_namespace:
+            raise ValueError(f"Marqo doesn't allow slashes in --index-namespace={args.index_namespace}.")
+
+    elif args.vector_store_provider == "pinecone":
+        if not os.getenv("PINECONE_API_KEY"):
+            raise ValueError("Please set the PINECONE_API_KEY environment variable.")
+        if not args.pinecone_index_name:
+            raise ValueError(f"Please set the vector_store.pinecone_index_name value.")
+
+
+def validate_indexing_args(args):
+    """Validates the indexing configuration and sets defaults."""
+    if args.include and args.exclude:
+        raise ValueError("At most one of indexing.include and indexing.exclude can be specified.")
+    if not args.include and not args.exclude:
+        args.exclude = pkg_resources.resource_filename(__name__, "sample-exclude.txt")
+    if args.include and not os.path.exists(args.include):
+        raise ValueError(f"Path --include={args.include} does not exist.")
+    if args.exclude and not os.path.exists(args.exclude):
+        raise ValueError(f"Path --exclude={args.exclude} does not exist.")
+    if not args.index_repo and not args.index_issues:
+        raise ValueError("Either --index_repo or --index_issues must be set to true.")
+    if args.index_issues and not os.getenv("GITHUB_TOKEN"):
+        raise ValueError("Please set the GITHUB_TOKEN environment variable.")

sage/configs/local.yaml

@@ -0,0 +1,16 @@
+# Embeddings
+embedding-provider: marqo
+embedding-model: hf/e5-base-v2
+tokens-per-chunk: 800
+chunks-per-batch: 64
+
+# Vector store
+vector-store-provider: marqo
+
+# LLM
+llm-provider: ollama
+llm-model: llama3.1
+
+# Reranking
+reranking-provider: huggingface
+reranking-model: cross-encoder/ms-marco-MiniLM-L-6-v2

sage/configs/remote.yaml

@@ -0,0 +1,18 @@
+# Embeddings
+embedding-provider: openai
+embedding-model: text-embedding-ada-002
+tokens-per-chunk: 800
+chunks-per-batch: 2000
+
+# Vector store
+vector-store-provider: pinecone
+pinecone-index-name: sage
+hybrid-retrieval: true
+
+# LLM
+llm-provider: openai
+llm-model: gpt-4
+
+# Reranking
+reranking-provider: cohere
+reranking-model: rerank-english-v3.0

sage/embedder.py

@@ -268,11 +268,11 @@ class MarqoEmbedder(BatchEmbedder):
 
 
 def build_batch_embedder_from_flags(data_manager: DataManager, chunker: Chunker, args) -> BatchEmbedder:
-    if args.embedder_type == "openai":
+    if args.embedding_provider == "openai":
         return OpenAIBatchEmbedder(data_manager, chunker, args.local_dir, args.embedding_model, args.embedding_size)
-    elif args.embedder_type == "marqo":
+    elif args.embedding_provider == "marqo":
         return MarqoEmbedder(
-            data_manager, chunker, index_name=args.index_name, url=args.marqo_url, model=args.embedding_model
+            data_manager, chunker, index_name=args.index_namespace, url=args.marqo_url, model=args.embedding_model
         )
     else:
-        raise ValueError(f"Unrecognized embedder type {args.embedder_type}")
+        raise ValueError(f"Unrecognized embedder type {args.embedding_provider}")

sage/index.py

@@ -1,196 +1,46 @@
 """Runs a batch job to compute embeddings for an entire repo and stores them into a vector store."""
 
-import argparse
 import logging
-import os
 import time
 
+import configargparse
 import pkg_resources
 
+import sage.config as sage_config
 from sage.chunker import UniversalFileChunker
 from sage.data_manager import GitHubRepoManager
 from sage.embedder import build_batch_embedder_from_flags
 from sage.github import GitHubIssuesChunker, GitHubIssuesManager
-from sage.vector_store import build_from_args
+from sage.vector_store import build_vector_store_from_args
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
 
-MARQO_MAX_CHUNKS_PER_BATCH = 64
-
-OPENAI_MAX_TOKENS_PER_CHUNK = 8192  # The ADA embedder from OpenAI has a maximum of 8192 tokens.
-OPENAI_MAX_CHUNKS_PER_BATCH = 2048  # The OpenAI batch embedding API enforces a maximum of 2048 chunks per batch.
-OPENAI_MAX_TOKENS_PER_JOB = (
-    3_000_000  # The OpenAI batch embedding API enforces a maximum of 3M tokens processed at once.
-)
-
-# Note that OpenAI embedding models have fixed dimensions, however, taking a slice of them is possible.
-# See "Reducing embedding dimensions" under https://platform.openai.com/docs/guides/embeddings/use-cases and
-# https://platform.openai.com/docs/api-reference/embeddings/create#embeddings-create-dimensions
-OPENAI_DEFAULT_EMBEDDING_SIZE = {
-    "text-embedding-ada-002": 1536,
-    "text-embedding-3-small": 1536,
-    "text-embedding-3-large": 3072,
-}
-
 
 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",
-        type=str,
-        default=None,
-        help="The embedding model. Defaults to `text-embedding-ada-002` for OpenAI and `hf/e5-base-v2` for Marqo.",
-    )
-    parser.add_argument(
-        "--embedding-size",
-        type=int,
-        default=None,
-        help="The embedding size to use for OpenAI text-embedding-3* models. Defaults to 1536 for small and 3072 for "
-        "large. Note that no other OpenAI models support a dynamic embedding size, nor do models used with Marqo.",
-    )
-    parser.add_argument("--vector-store-type", default="marqo", choices=["pinecone", "marqo"])
-    parser.add_argument(
-        "--local-dir",
-        default="repos",
-        help="The local directory to store the repository",
-    )
-    parser.add_argument(
-        "--tokens-per-chunk",
-        type=int,
-        default=800,
-        help="https://arxiv.org/pdf/2406.14497 recommends a value between 200-800.",
-    )
-    parser.add_argument(
-        "--chunks-per-batch",
-        type=int,
-        help="Maximum chunks per batch. We recommend 2000 for the OpenAI embedder. Marqo enforces a limit of 64.",
-    )
-    parser.add_argument(
-        "--pinecone-index-name",
-        default=None,
-        help="Pinecone index name. Required if using Pinecone as the vector store. If the index doesn't exist already, "
-        "we will create it.",
-    )
-    parser.add_argument(
-        "--index-namespace",
-        default=None,
-        help="Index namespace for this repo. When not specified, we default it to a derivative of the repo name."
-    )
-    parser.add_argument(
-        "--include",
-        help="Path to a file containing a list of extensions to include. One extension per line.",
-    )
-    parser.add_argument(
-        "--exclude",
-        help="Path to a file containing a list of extensions to exclude. One extension per line.",
-    )
-    parser.add_argument(
-        "--max-embedding-jobs",
-        type=int,
-        help="Maximum number of embedding jobs to run. Specifying this might result in "
-        "indexing only part of the repository, but prevents you from burning through OpenAI credits.",
+    parser = configargparse.ArgParser(
+        description="Batch-embeds a GitHub repository and its issues.", ignore_unknown_config_file_keys=True
     )
-    parser.add_argument(
-        "--marqo-url",
-        default="http://localhost:8882",
-        help="URL for the Marqo server. Required if using Marqo as embedder or vector store.",
-    )
-    # Pass --no-index-repo in order to not index the repository.
-    parser.add_argument(
-        "--index-repo",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Whether to index the repository. At least one of --index-repo and --index-issues must be True.",
-    )
-    # Pass --no-index-issues in order to not index the issues.
-    parser.add_argument(
-        "--index-issues",
-        action=argparse.BooleanOptionalAction,
-        default=False,
-        help="Whether to index GitHub issues. At least one of --index-repo and --index-issues must be True. When "
-        "--index-issues is set, you must also set a GITHUB_TOKEN environment variable.",
-    )
-    parser.add_argument(
-        "--index-issue-comments",
-        action=argparse.BooleanOptionalAction,
-        default=False,
-        help="Whether to index the comments of GitHub issues. This is only relevant if --index-issues is set. "
-        "GitHub's API for downloading comments is quite slow. Indexing solely the body of an issue seems to bring most "
-        "of the gains anyway.",
-    )
-    parser.add_argument(
-        "--hybrid-retrieval",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Whether to use a hybrid of vector DB + BM25 retrieval. When set to False, we only use vector DB "
-        "retrieval. This is only relevant if using Pinecone as the vector store.",
-    )
-    args = parser.parse_args()
+    sage_config.add_config_args(parser)
 
-    # Validate embedder and vector store compatibility.
-    if args.embedder_type == "openai" and args.vector_store_type != "pinecone":
-        parser.error("When using OpenAI embedder, the vector store type must be Pinecone.")
-    if args.embedder_type == "marqo" and args.vector_store_type != "marqo":
-        parser.error("When using the marqo embedder, the vector store type must also be marqo.")
-    if args.vector_store_type == "marqo":
-        if "/" in args.index_namespace:
-            parser.error("The index namespace cannot contain slashes when using Marqo as the vector store.")
-    elif args.vector_store_type == "pinecone" and not args.pinecone_index_name:
-        parser.error("When using Pinecone as the vector store, you must specify --pinecone-index-name")
+    arg_validators = [
+        sage_config.add_repo_args(parser),
+        sage_config.add_embedding_args(parser),
+        sage_config.add_vector_store_args(parser),
+        sage_config.add_indexing_args(parser),
+    ]
 
-    # Validate embedder parameters.
-    if args.embedder_type == "marqo":
-        if args.embedding_model is None:
-            args.embedding_model = "hf/e5-base-v2"
-        if args.chunks_per_batch is None:
-            args.chunks_per_batch = MARQO_MAX_CHUNKS_PER_BATCH
-        elif args.chunks_per_batch > MARQO_MAX_CHUNKS_PER_BATCH:
-            args.chunks_per_batch = MARQO_MAX_CHUNKS_PER_BATCH
-            logging.warning(
-                f"Marqo enforces a limit of {MARQO_MAX_CHUNKS_PER_BATCH} chunks per batch. "
-                "Overwriting --chunks_per_batch."
-            )
-    elif args.embedder_type == "openai":
-        if args.tokens_per_chunk > OPENAI_MAX_TOKENS_PER_CHUNK:
-            args.tokens_per_chunk = OPENAI_MAX_TOKENS_PER_CHUNK
-            logging.warning(
-                f"OpenAI enforces a limit of {OPENAI_MAX_TOKENS_PER_CHUNK} tokens per chunk. "
-                "Overwriting --tokens_per_chunk."
-            )
-        if args.chunks_per_batch is None:
-            args.chunks_per_batch = 2000
-        elif args.chunks_per_batch > OPENAI_MAX_CHUNKS_PER_BATCH:
-            args.chunks_per_batch = OPENAI_MAX_CHUNKS_PER_BATCH
-            logging.warning(
-                f"OpenAI enforces a limit of {OPENAI_MAX_CHUNKS_PER_BATCH} chunks per batch. "
-                "Overwriting --chunks_per_batch."
-            )
-        if args.tokens_per_chunk * args.chunks_per_batch >= OPENAI_MAX_TOKENS_PER_JOB:
-            parser.error(f"The maximum number of chunks per job is {OPENAI_MAX_TOKENS_PER_JOB}.")
-        if args.embedding_model is None:
-            args.embedding_model = "text-embedding-ada-002"
-        if args.embedding_size is None:
-            args.embedding_size = OPENAI_DEFAULT_EMBEDDING_SIZE.get(args.embedding_model)
+    args = parser.parse_args()
 
-    if args.include and args.exclude:
-        parser.error("At most one of --include and --exclude can be specified.")
-    if not args.include and not args.exclude:
-        args.exclude = pkg_resources.resource_filename(__name__, "sample-exclude.txt")
-    if not args.index_repo and not args.index_issues:
-        parser.error("At least one of --index-repo and --index-issues must be true.")
+    for validator in arg_validators:
+        validator(args)
 
-    # Fail early on missing environment variables.
-    if args.embedder_type == "openai" and not os.getenv("OPENAI_API_KEY"):
-        parser.error("Please set the OPENAI_API_KEY environment variable.")
-    if args.vector_store_type == "pinecone" and not os.getenv("PINECONE_API_KEY"):
-        parser.error("Please set the PINECONE_API_KEY environment variable.")
-    if args.index_issues and not os.getenv("GITHUB_TOKEN"):
-        parser.error("Please set the GITHUB_TOKEN environment variable.")
+    # 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.")
+    if args.embedding_provider == "marqo" and args.vector_store_provider != "marqo":
+        parser.error("When using the marqo embedder, the vector store type must also be marqo.")
 
     ######################
     # Step 1: Embeddings #
@@ -228,7 +78,7 @@ def main():
     # Step 2: Vector Store #
     ########################
 
-    if args.vector_store_type == "marqo":
+    if args.vector_store_provider == "marqo":
         # Marqo computes embeddings and stores them in the vector store at once, so we're done.
         logging.info("Done!")
         return
@@ -240,7 +90,7 @@ def main():
             time.sleep(30)
 
         logging.info("Moving embeddings to the repo vector store...")
-        repo_vector_store = build_from_args(args)
+        repo_vector_store = build_vector_store_from_args(args)
         repo_vector_store.ensure_exists()
         repo_vector_store.upsert(repo_embedder.download_embeddings(repo_jobs_file))
 
@@ -251,7 +101,7 @@ def main():
             time.sleep(30)
 
         logging.info("Moving embeddings to the issues vector store...")
-        issues_vector_store = build_from_args(args)
+        issues_vector_store = build_vector_store_from_args(args)
         issues_vector_store.ensure_exists()
         issues_vector_store.upsert(issues_embedder.download_embeddings(issues_jobs_file))
 

sage/llm.py

@@ -10,12 +10,12 @@ def build_llm_via_langchain(provider: str, model: str):
     if provider == "openai":
         if "OPENAI_API_KEY" not in os.environ:
             raise ValueError("Please set the OPENAI_API_KEY environment variable.")
-        return ChatOpenAI(model=model)
+        return ChatOpenAI(model=model or "gpt-4")
     elif provider == "anthropic":
         if "ANTHROPIC_API_KEY" not in os.environ:
             raise ValueError("Please set the ANTHROPIC_API_KEY environment variable.")
-        return ChatAnthropic(model=model)
+        return ChatAnthropic(model=model or "claude-3-opus-20240229")
     elif provider == "ollama":
-        return ChatOllama(model=model)
+        return ChatOllama(model=model or "llama3.1")
     else:
         raise ValueError(f"Unrecognized LLM provider {provider}. Contributons are welcome!")

sage/vector_store.py

@@ -146,35 +146,16 @@ class MarqoVectorStore(VectorStore):
         return vectorstore.as_retriever(search_kwargs={"k": top_k})
 
 
-def build_from_args(args: dict) -> VectorStore:
+def build_vector_store_from_args(args: dict) -> VectorStore:
     """Builds a vector store from the given command-line arguments."""
-    if args.vector_store_type == "pinecone":
-        if not args.pinecone_index_name:
-            raise ValueError("Please specify --pinecone-index-name for Pinecone.")
-        dimension = args.embedding_size if "embedding_size" in args else None
-
-        index_namespace = args.index_namespace
-        if not index_namespace:
-            index_namespace = args.repo_id
-            if args.commit_hash:
-                namespace += "/" + args.commit_hash
-
+    if args.vector_store_provider == "pinecone":
         return PineconeVectorStore(
             index_name=args.pinecone_index_name,
-            namespace=index_namespace,
-            dimension=dimension,
+            namespace=args.index_namespace,
+            dimension=args.embedding_size if "embedding_size" in args else None,
             hybrid=args.hybrid_retrieval,
         )
-    elif args.vector_store_type == "marqo":
-        marqo_url = args.marqo_url or "http://localhost:8882"
-
-        index_namespace = args.index_namespace
-        if not index_namespace:
-            # Marqo doesn't allow slashes in the index name.
-            index_namespace = args.repo_id.split("/")[1]
-            if args.commit_hash:
-                index_namespace += "_" + args.commit_hash
-
-        return MarqoVectorStore(url=marqo_url, index_name=index_namespace)
+    elif args.vector_store_provider == "marqo":
+        return MarqoVectorStore(url=args.marqo_url, index_name=args.index_namespace)
     else:
-        raise ValueError(f"Unrecognized vector store type {args.vector_store_type}")
+        raise ValueError(f"Unrecognized vector store type {args.vector_store_provider}")