Add official documentation instead of README (#93)

* Add Anthropic reranker. Switch to "remote" config by default.

* Update README to point to docs.

* Format tests.

README.md

@@ -8,6 +8,11 @@
         <a href="https://github.com/Storia-AI/sage/stargazers" target=="_blank"><img alt="GitHub Repo stars" src="https://img.shields.io/github/stars/Storia-AI/sage?logo=github&link=https%3A%2F%2Fgithub.com%2FStoria-AI%2Fsage%2Fstargazers"></a>
         <a href="https://github.com/Storia-AI/sage/blob/main/LICENSE" target=="_blank"><img alt="GitHub License" src="https://img.shields.io/github/license/Storia-AI/sage" /></a>
     </div>
+    <div>
+        <a href="https://sage-docs.storia.ai">Documentation</a>
+        <span>&#183;</span>
+        <a href="https://sage.storia.ai">Hosted app</a>
+    </div>
     <br />
     <figure>
         <!-- The <kbd> and <sub> tags are work-arounds for styling, since GitHub doesn't take into account inline styles. Note it might display awkwardly on other Markdown editors. -->
@@ -16,253 +21,19 @@
     </figure>
 </div>
 
-# Getting started
-
-## Installation
-
-<details open>
-<summary><strong>Using pipx (for regular users) </strong></summary>
-Make sure pipx is installed on your system (see <a href="https://pipx.pypa.io/stable/installation/">instructions</a>), then run:
-
-```
-pipx install git+https://github.com/Storia-AI/sage.git@main
-```
-
-</details>
-
-<details>
-<summary><strong>Using venv and pip (for contributors)</strong></summary>
-Alternatively, you can manually create a virtual environment and install Code Sage via pip:
-
-```
-python -m venv sage-venv
-source sage-venv/bin/activate
-git clone https://github.com/Storia-AI/sage.git
-cd sage
-pip install -e .
-```
-
-</details>
-
-## Prerequisites
-
-`sage` performs two steps:
-
-1. Indexes your codebase (requiring an embedder and a vector store)
-2. Enables chatting via LLM + RAG (requiring access to an LLM)
-
-<details open>
-<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:
-
-    ```
-    docker rm -f marqo
-    docker pull marqoai/marqo:latest
-    docker run --name marqo -it -p 8882:8882 marqoai/marqo:latest
-    ```
-
-    This will open a persistent Marqo console window. This should take around 2-3 minutes on a fresh install.
-
-2. To chat with an LLM locally, we use <a href="https://github.com/ollama/ollama">Ollama</a>:
-
-    - Head over to [ollama.com](https://ollama.com) to download the appropriate binary for your machine.
-    - Open a new terminal window
-    - Pull the desired model, e.g. `ollama pull llama3.1`.
-
-</details>
-
-<details>
-<summary><strong>:cloud: Using external providers (higher quality)</strong></summary>
-
-1. For embeddings, we support <a href="https://platform.openai.com/docs/guides/embeddings">OpenAI</a> and <a href="https://docs.voyageai.com/docs/embeddings">Voyage</a>. According to [our experiments](benchmarks/retrieval/README.md), OpenAI is better quality. Their batch API is also faster, with more generous rate limits. Export the API key of the desired provider:
-
-    ```
-    export OPENAI_API_KEY=... # or
-    export VOYAGE_API_KEY=...
-    ```
-
-2. We use <a href="https://www.pinecone.io/">Pinecone</a> for the vector store, so you will need an API key:
-
-    ```
-    export PINECONE_API_KEY=...
-    ```
-    If you want to reuse an existing Pinecone index, specify it. Otherwise we'll create a new one called `sage`.
-    ```
-    export PINECONE_INDEX_NAME=...
-    ```
-
-3. For reranking, we support <a href="https://developer.nvidia.com/blog/enhancing-rag-pipelines-with-re-ranking/">NVIDIA</a>, <a href="https://docs.voyageai.com/docs/reranker">Voyage</a>, <a href="https://cohere.com/rerank">Cohere</a>, and <a href="https://jina.ai/reranker/">Jina</a>.
-   - According to [our experiments](benchmark/retrieval/README.md), NVIDIA performs best. To get an API key, follow [these instructions](https://docs.nvidia.com/nim/large-language-models/latest/getting-started.html#generate-an-api-key). Note that NVIDIA's API keys are model-specific. We recommend using `nvidia/nv-rerankqa-mistral-4b-v3`.
-   - Export the API key of the desired provider:
-    ```
-    export NVIDIA_API_KEY=...  # or
-    export VOYAGE_API_KEY=...  # or
-    export COHERE_API_KEY=...  # or
-    export JINA_API_KEY=...
-    ```
-
-4. For chatting with an LLM, we support OpenAI and Anthropic. For the latter, set an additional API key:
-    ```
-    export ANTHROPIC_API_KEY=...
-    ```
-
-For easier configuration, adapt the entries within the sample `.sage-env` (change the API keys names based on your desired setup) and run:
-```
-source .sage-env
-```
-</details>
-
-### 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
-
-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
-    ```
-    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 use external providers instead of running locally, set `--mode=remote`.
-</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
-
-<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>
-
-You can specify an inclusion or exclusion file in the following format:
-```
-# This is a comment
-ext:.my-ext-1
-ext:.my-ext-2
-ext:.my-ext-3
-dir:my-dir-1
-dir:my-dir-2
-dir:my-dir-3
-file:my-file-1.md
-file:my-file-2.py
-file:my-file-3.cpp
-```
-where:
-- `ext` specifies a file extension
-- `dir` specifies a directory. This is not a full path. For instance, if you specify `dir:tests` in an exclusion directory, then a file like `/path/to/my/tests/file.py` will be ignored.
-- `file` specifies a file name. This is also not a full path. For instance, if you specify `file:__init__.py`, then a file like `/path/to/my/__init__.py` will be ignored.
-
-To specify an inclusion file (i.e. only index the specified files):
-```
-sage-index $GITHUB_REPO --include=/path/to/inclusion/file
-```
-
-To specify an exclusion file (i.e. index all files, except for the ones specified):
-```
-sage-index $GITHUB_REPO --exclude=/path/to/exclusion/file
-```
-By default, we use the exclusion file [sample-exclude.txt](sage/sample-exclude.txt).
-
-</details>
-
-<details>
-<summary><strong>:bug: Index open GitHub issues</strong></summary>
-
-You will need a GitHub token first:
-
-```
-export GITHUB_TOKEN=...
-```
-
-To index GitHub issues without comments:
-```
-sage-index $GITHUB_REPO --index-issues
-```
-
-To index GitHub issues with comments:
-```
-sage-index $GITHUB_REPO --index-issues --index-issue-comments
-```
-
-To index GitHub issues, but not the codebase:
-```
-sage-index $GITHUB_REPO --index-issues --no-index-repo
-```
-
-</details>
-
-<details>
-<summary><strong>:books: Experiment with retrieval strategies</strong></summary>
-
-Retrieving the right files from the vector database is arguably the quality bottleneck of the system. We are actively experimenting with various retrieval strategies and documenting our findings [here](benchmark/retrieval/README.md).
-
-Currently, we support the following types of retrieval:
-- **Vanilla RAG** from a vector database (nearest neighbor between dense embeddings). This is the default.
-- **Hybrid RAG** that combines dense retrieval (embeddings-based) with sparse retrieval (BM25). Use `--retrieval-alpha` to weigh the two strategies.
-
-    - A value of 1 means dense-only retrieval and 0 means BM25-only retrieval.
-    - Note this is not available when running locally, only when using Pinecone as a vector store.
-    - Contrary to [Anthropic's findings](https://www.anthropic.com/news/contextual-retrieval), we find that BM25 is actually damaging performance *on codebases*, because it gives undeserved advantage to Markdown files.
-
-- **Multi-query retrieval** performs multiple query rewrites, makes a separate retrieval call for each, and takes the union of the retrieved documents. You can activate it by passing `--multi-query-retrieval`. This can be combined with both vanilla and hybrid RAG.
-
-    - We find that [on our benchmark](benchmark/retrieval/README.md) this only marginally improves retrieval quality (from 0.44 to 0.46 R-precision) while being significantly slower and more expensive due to LLM calls. But your mileage may vary.
-
-- **LLM-only retrieval** completely circumvents indexing the codebase. We simply enumerate all file paths and pass them to an LLM together with the user query. We ask the LLM which files are likely to be relevant for the user query, solely based on their filenames. You can activate it by passing `--llm-retriever`.
-
-    - We find that [on our benchmark](benchmark/retrieval/README.md) the performance is comparable with vector database solutions (R-precision is 0.44 for both). This is quite remarkable, since we've saved so much effort by not indexing the codebase. However, we are reluctant to claim that these findings generalize, for the following reasons:
-        - Our (artificial) dataset occasionally contains explicit path names in the query, making it trivial for the LLM. Sample query: *"Alice is managing a series of machine learning experiments. Please explain in detail how `main` in `examples/pytorch/image-pretraining/run_mim.py` allows her to organize the outputs of each experiment in separate directories."*
-        - Our benchmark focuses on the Transformers library, which is well-maintained and the file paths are often meaningful. This might not be the case for all codebases.
-
-</details>
-
-# Why chat with a codebase?
-
-Sometimes you just want to learn how a codebase works and how to integrate it, without spending hours sifting through
-the code itself.
-
-`sage` is like an open-source GitHub Copilot with the most up-to-date information about your repo.
-
-Features:
-
-- **Dead-simple set-up.** Run *two scripts* and you have a functional chat interface for your code. That's really it.
-- **Heavily documented answers.** Every response shows where in the code the context for the answer was pulled from. Let's build trust in the AI.
-- **Runs locally or on the cloud.**
-- **Plug-and-play.** Want to improve the algorithms powering the code understanding/generation? We've made every component of the pipeline easily swappable. Google-grade engineering standards allow you to customize to your heart's content.
+***
 
-# Changelog
+**Sage** is like an open-source GitHub Copilot that helps you learn how a codebase works and how to integrate it into your project without spending hours sifting through the code.
 
-- 2024-09-16: Renamed `repo2vec` to `sage`.
-- 2024-09-03: Support for indexing GitHub issues.
-- 2024-08-30: Support for running everything locally (Marqo for embeddings, Ollama for LLMs).
+# Main features
+- **Dead-simple setup**. Follow our [quickstart guide](https://sage-docs.storia.ai/quickstart) to get started.
+- **Runs locally or on the cloud**. When privacy is your priority, you can run the entire pipeline locally using [Ollama](https://ollama.com) for LLMs and [Marqo](https://github.com/marqo-ai/marqo) as a vector store. When optimizing for quality, you can use third-party LLM providers like OpenAI and Anthropic.
+- **Wide range of built-in retrieval mechanisms**. We support both lightweight retrieval strategies (with nothing more but an LLM API key required) and more traditional RAG (which requires indexing the codebase). There are many knobs you can tune for retrieval to work well on your codebase.
+- **Well-documented experiments**. We profile various strategies (for embeddings, retrieval etc.) on our own benchmark and thoroughly [document the results](benchmarks/retrieval/README.md).
 
 # Want your repository hosted?
 
-We're working to make all code on the internet searchable and understandable for devs. You can check out our early product, [Code Sage](https://sage.storia.ai). We pre-indexed a slew of OSS repos, and you can index your desired ones by simply pasting a GitHub URL.
+We're working to make all code on the internet searchable and understandable for devs. You can check out [hosted app](https://sage.storia.ai). We pre-indexed a slew of OSS repos, and you can index your desired ones by simply pasting a GitHub URL.
 
 If you're the maintainer of an OSS repo and would like a dedicated page on Code Sage (e.g. `sage.storia.ai/your-repo`), then send us a message at [founders@storia.ai](mailto:founders@storia.ai). We'll do it for free!
 

sage/config.py

@@ -58,8 +58,8 @@ def add_config_args(parser: ArgumentParser):
     parser.add(
         "--mode",
         choices=["local", "remote"],
-        default="local",
-        help="Whether to use local-only resources or call third-party providers.",
+        default="remote",
+        help="Whether to use local-only resources or call third-party providers (remote).",
     )
     parser.add(
         "--config",

sage/configs/remote.yaml

@@ -1,17 +1,15 @@
+llm-retriever: true
+llm-provider: anthropic
+reranker-provider: anthropic
+
+# The settings below (embeddings and vector store) are only relevant when setting --no-llm-retriever
+
 # Embeddings
 embedding-provider: openai
 embedding-model: text-embedding-3-small
 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-4o
-
-# Reranking
-reranker-provider: nvidia

sage/reranker.py

@@ -1,14 +1,21 @@
+import logging
 import os
 from enum import Enum
-from typing import Optional
+from typing import List, Optional
 
 from langchain.retrievers.document_compressors import CrossEncoderReranker
 from langchain_cohere import CohereRerank
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 from langchain_community.document_compressors import JinaRerank
-from langchain_core.documents import BaseDocumentCompressor
+from langchain_core.callbacks.manager import Callbacks
+from langchain_core.documents import BaseDocumentCompressor, Document
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.prompts import PromptTemplate
 from langchain_nvidia_ai_endpoints import NVIDIARerank
 from langchain_voyageai import VoyageAIRerank
+from pydantic import ConfigDict, Field
+
+from sage.llm import build_llm_via_langchain
 
 
 class RerankerProvider(Enum):
@@ -18,6 +25,58 @@ class RerankerProvider(Enum):
     NVIDIA = "nvidia"
     JINA = "jina"
     VOYAGE = "voyage"
+    # Anthropic doesn't provide an explicit reranker; we simply prompt the LLM with the user query and the content of
+    # the top k documents.
+    ANTHROPIC = "anthropic"
+
+
+class LLMReranker(BaseDocumentCompressor):
+    """Reranker that passes the user query and top N documents to a language model to order them.
+
+    Note that Langchain's RerankLLM does not support LLMs from Anthropic.
+    https://python.langchain.com/api_reference/community/document_compressors/langchain_community.document_compressors.rankllm_rerank.RankLLMRerank.html
+    Also, they rely on https://github.com/castorini/rank_llm, which doesn't run on Apple Silicon (M1/M2 chips).
+    """
+
+    llm: BaseLanguageModel = Field(...)
+    top_k: int = Field(...)
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="forbid",
+    )
+
+    @property
+    def prompt(self):
+        return PromptTemplate.from_template(
+            "Given the following query: '{query}'\n\n"
+            "And these documents:\n\n{documents}\n\n"
+            "Rank the documents based on their relevance to the query. "
+            "Return only the document numbers in order of relevance, separated by commas. For example: 2,5,1,3,4. "
+            "Return absolutely nothing else."
+        )
+
+    def compress_documents(
+        self,
+        documents: List[Document],
+        query: str,
+        callbacks: Optional[Callbacks] = None,
+    ) -> List[Document]:
+        if len(documents) <= self.top_k:
+            return documents
+
+        doc_texts = [f"Document {i+1}:\n{doc.page_content}\n" for i, doc in enumerate(documents)]
+        docs_str = "\n".join(doc_texts)
+
+        llm_input = self.prompt.format(query=query, documents=docs_str)
+        result = self.llm.predict(llm_input)
+
+        try:
+            ranked_indices = [int(idx) - 1 for idx in result.strip().split(",")][: self.top_k]
+            return [documents[i] for i in ranked_indices]
+        except ValueError:
+            logging.warning("Failed to parse reranker output. Returning original order. LLM responded with: %s", result)
+            return documents[: self.top_k]
 
 
 def build_reranker(provider: str, model: Optional[str] = None, top_k: Optional[int] = 5) -> BaseDocumentCompressor:
@@ -46,4 +105,7 @@ def build_reranker(provider: str, model: Optional[str] = None, top_k: Optional[i
             raise ValueError("Please set the VOYAGE_API_KEY environment variable")
         model = model or "rerank-1"
         return VoyageAIRerank(model=model, api_key=os.environ.get("VOYAGE_API_KEY"), top_k=top_k)
+    if provider == RerankerProvider.ANTHROPIC.value:
+        llm = build_llm_via_langchain("anthropic", model)
+        return LLMReranker(llm=llm, top_k=1)
     raise ValueError(f"Invalid reranker provider: {provider}")

tests/test_data_manager.py

@@ -14,102 +14,93 @@ Here, one can find a set of unit tests for GitHubRepoManager class from sage.dat
 
 """
 
-
+import os
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
+
 from sage.data_manager import GitHubRepoManager
-import os
+
+
 class TestGitHubRepoManager(unittest.TestCase):
-    @patch('git.Repo.clone_from')
+    @patch("git.Repo.clone_from")
     def test_download_clone_success(self, mock_clone):
         """Test the download() method of GitHubRepoManager by mocking the cloning process."""
-        repo_manager = GitHubRepoManager(
-            repo_id='Storia-AI/sage',
-            local_dir='/tmp/test_repo'
-        )
-        mock_clone.return_value = MagicMock() 
+        repo_manager = GitHubRepoManager(repo_id="Storia-AI/sage", local_dir="/tmp/test_repo")
+        mock_clone.return_value = MagicMock()
         result = repo_manager.download()
         mock_clone.assert_called_once_with(
-            'https://github.com/Storia-AI/sage.git',
-            '/tmp/test_repo/Storia-AI/sage',
-            depth=1,
-            single_branch=True
+            "https://github.com/Storia-AI/sage.git", "/tmp/test_repo/Storia-AI/sage", depth=1, single_branch=True
         )
         self.assertTrue(result)
 
-    @patch('sage.data_manager.requests.get')
+    @patch("sage.data_manager.requests.get")
     def test_is_public_repository(self, mock_get):
         """Test the is_public property to check if a repository is public."""
         mock_get.return_value.status_code = 200
-        repo_manager = GitHubRepoManager(repo_id='Storia-AI/sage')
+        repo_manager = GitHubRepoManager(repo_id="Storia-AI/sage")
         self.assertTrue(repo_manager.is_public)
-        mock_get.assert_called_once_with('https://api.github.com/repos/Storia-AI/sage', timeout=10)
+        mock_get.assert_called_once_with("https://api.github.com/repos/Storia-AI/sage", timeout=10)
 
-    @patch('sage.data_manager.requests.get')
+    @patch("sage.data_manager.requests.get")
     def test_is_private_repository(self, mock_get):
         """Test the is_public property to check if a repository is private."""
-        mock_get.return_value.status_code = 404 
-        repo_manager = GitHubRepoManager(repo_id='Storia-AI/sage')
+        mock_get.return_value.status_code = 404
+        repo_manager = GitHubRepoManager(repo_id="Storia-AI/sage")
         self.assertFalse(repo_manager.is_public)
-        mock_get.assert_called_once_with('https://api.github.com/repos/Storia-AI/sage', timeout=10)
+        mock_get.assert_called_once_with("https://api.github.com/repos/Storia-AI/sage", timeout=10)
 
-    @patch('sage.data_manager.requests.get')
+    @patch("sage.data_manager.requests.get")
     def test_default_branch(self, mock_get):
         """Test the default_branch property to fetch the default branch of the repository."""
         mock_get.return_value.status_code = 200
-        mock_get.return_value.json.return_value = {'default_branch': 'main'}
-        repo_manager = GitHubRepoManager(repo_id='Storia-AI/sage')
-        self.assertEqual(repo_manager.default_branch, 'main')
+        mock_get.return_value.json.return_value = {"default_branch": "main"}
+        repo_manager = GitHubRepoManager(repo_id="Storia-AI/sage")
+        self.assertEqual(repo_manager.default_branch, "main")
         mock_get.assert_called_once_with(
-            'https://api.github.com/repos/Storia-AI/sage',
-            headers={'Accept': 'application/vnd.github.v3+json'}
+            "https://api.github.com/repos/Storia-AI/sage", headers={"Accept": "application/vnd.github.v3+json"}
         )
 
-    @patch('builtins.open', new_callable=unittest.mock.mock_open, read_data='ext:.py\nfile:test.py\ndir:test_dir\n')
+    @patch("builtins.open", new_callable=unittest.mock.mock_open, read_data="ext:.py\nfile:test.py\ndir:test_dir\n")
     def test_parse_filter_file(self, mock_file):
         """Test the _parse_filter_file method for correct parsing of inclusion/exclusion files."""
-        repo_manager = GitHubRepoManager(repo_id='Storia-AI/sage', inclusion_file='dummy_path')
-        expected = {
-            'ext': ['.py'],
-            'file': ['test.py'],
-            'dir': ['test_dir']
-        }
-        result = repo_manager._parse_filter_file('dummy_path')
+        repo_manager = GitHubRepoManager(repo_id="Storia-AI/sage", inclusion_file="dummy_path")
+        expected = {"ext": [".py"], "file": ["test.py"], "dir": ["test_dir"]}
+        result = repo_manager._parse_filter_file("dummy_path")
         self.assertEqual(result, expected)
 
-    @patch('os.path.exists')
-    @patch('os.remove')  
-    @patch('builtins.open', new_callable=unittest.mock.mock_open, read_data='dummy content')
+    @patch("os.path.exists")
+    @patch("os.remove")
+    @patch("builtins.open", new_callable=unittest.mock.mock_open, read_data="dummy content")
     def test_walk_included_files(self, mock_open, mock_remove, mock_exists):
         """Test the walk method to ensure it only includes specified files."""
         mock_exists.return_value = True
-        repo_manager = GitHubRepoManager(
-            repo_id="Storia-AI/sage",
-            local_dir="/tmp/test_repo"
-        )
-        with patch('os.walk', return_value=[
-            ('/tmp/test_repo', ('subdir',), ('included_file.py', 'excluded_file.txt')),
-        ]):
+        repo_manager = GitHubRepoManager(repo_id="Storia-AI/sage", local_dir="/tmp/test_repo")
+        with patch(
+            "os.walk",
+            return_value=[
+                ("/tmp/test_repo", ("subdir",), ("included_file.py", "excluded_file.txt")),
+            ],
+        ):
             included_files = list(repo_manager.walk())
             print("Included files:", included_files)
-            self.assertTrue(any(file[1]['file_path'] == 'included_file.py' for file in included_files))
+            self.assertTrue(any(file[1]["file_path"] == "included_file.py" for file in included_files))
+
     def test_read_file(self):
         """Test the read_file method to read the content of a file."""
-        mock_file_path = '/tmp/test_repo/test_file.txt'
-        with patch('builtins.open', new_callable=unittest.mock.mock_open, read_data='Hello, World!'):
-            repo_manager = GitHubRepoManager(repo_id='Storia-AI/sage', local_dir='/tmp/test_repo')
-            content = repo_manager.read_file('test_file.txt')
-            self.assertEqual(content, 'Hello, World!')
+        mock_file_path = "/tmp/test_repo/test_file.txt"
+        with patch("builtins.open", new_callable=unittest.mock.mock_open, read_data="Hello, World!"):
+            repo_manager = GitHubRepoManager(repo_id="Storia-AI/sage", local_dir="/tmp/test_repo")
+            content = repo_manager.read_file("test_file.txt")
+            self.assertEqual(content, "Hello, World!")
 
-    @patch('os.makedirs')
+    @patch("os.makedirs")
     def test_create_log_directories(self, mock_makedirs):
         """Test that log directories are created."""
-        repo_manager = GitHubRepoManager(
-            repo_id="Storia-AI/sage",
-            local_dir="/tmp/test_repo"
-        )
+        repo_manager = GitHubRepoManager(repo_id="Storia-AI/sage", local_dir="/tmp/test_repo")
 
         with self.assertRaises(AttributeError):
             repo_manager.create_log_directories()
-if __name__ == 'main':
+
+
+if __name__ == "main":
     unittest.main()

tests/test_github.py

@@ -9,13 +9,22 @@ pip install pytest
 pip install pytest-mock
 """
 
+from unittest.mock import MagicMock, patch
+
 import pytest
 import requests
-from unittest.mock import patch, MagicMock
-from sage.github import GitHubIssuesManager, GitHubIssue, GitHubIssueComment, GitHubIssuesChunker, IssueChunk
+
+from sage.github import (
+    GitHubIssue,
+    GitHubIssueComment,
+    GitHubIssuesChunker,
+    GitHubIssuesManager,
+    IssueChunk,
+)
+
 
 class TestGitHubIssuesManager:
-    
+
     @pytest.fixture(autouse=True)
     def setup_method(self):
         """Fixture to create a GitHubIssuesManager instance for each test."""
@@ -24,33 +33,37 @@ class TestGitHubIssuesManager:
     @staticmethod
     def mock_issue_response():
         """A mock response for GitHub issues."""
-        return MagicMock(json=lambda: [
-            {
-                "url": "https://api.github.com/repos/random/random-repo/issues/1",
-                "html_url": "https://github.com/random/random-repo/issues/1",
-                "title": "Found a bug", 
-                "body": "I'm having a problem with this.",
-                "comments_url": "https://api.github.com/repos/random/random-repo/issues/1/comments",
-                "comments": 2
-            }
-        ])
-    
+        return MagicMock(
+            json=lambda: [
+                {
+                    "url": "https://api.github.com/repos/random/random-repo/issues/1",
+                    "html_url": "https://github.com/random/random-repo/issues/1",
+                    "title": "Found a bug",
+                    "body": "I'm having a problem with this.",
+                    "comments_url": "https://api.github.com/repos/random/random-repo/issues/1/comments",
+                    "comments": 2,
+                }
+            ]
+        )
+
     @staticmethod
     def mock_comment_response():
         """Create a mock response for GitHub issue comments."""
-        return MagicMock(json=lambda: [
-            {
-                "url": "https://api.github.com/repos/random/random-repo/issues/comments/1",
-                "html_url": "https://github.com/random/random-repo/issues/comments/1",
-                "body": "This is a comment."
-            }
-        ])
+        return MagicMock(
+            json=lambda: [
+                {
+                    "url": "https://api.github.com/repos/random/random-repo/issues/comments/1",
+                    "html_url": "https://github.com/random/random-repo/issues/comments/1",
+                    "body": "This is a comment.",
+                }
+            ]
+        )
 
-    @patch('github.requests.get')
+    @patch("github.requests.get")
     def test_download_issues(self, mock_get):
         """Test the download of issues from GitHub."""
         mock_get.side_effect = [self.mock_issue_response(), self.mock_comment_response()]
-        
+
         self.github_manager.download()
 
         assert len(self.github_manager.issues) == 1
@@ -58,21 +71,27 @@ class TestGitHubIssuesManager:
         assert self.github_manager.issues[0].body == "I'm having a problem with this."
         assert self.github_manager.issues[0].url == "https://api.github.com/repos/random/random-repo/issues/1"
 
-    @patch('github.requests.get')
+    @patch("github.requests.get")
     def test_walk_issues(self, mock_get):
         """Test the walking through downloaded issues."""
         self.github_manager.issues = [
             GitHubIssue(url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[]),
-            GitHubIssue(url="issue_url_2", html_url="html_issue_url_2", title="Another Test Issue", body="Another Test Body", comments=[]),
+            GitHubIssue(
+                url="issue_url_2",
+                html_url="html_issue_url_2",
+                title="Another Test Issue",
+                body="Another Test Body",
+                comments=[],
+            ),
         ]
 
         issues = list(self.github_manager.walk())
-        
+
         assert len(issues) == 2
         assert issues[0][0].title == "Test Issue"
         assert issues[1][0].title == "Another Test Issue"
 
-    @patch('github.requests.get')
+    @patch("github.requests.get")
     def test_get_page_of_issues(self, mock_get):
         """Test fetching a page of issues."""
         mock_response = MagicMock()
@@ -80,23 +99,25 @@ class TestGitHubIssuesManager:
             {
                 "url": "https://api.github.com/repos/random/random-repo/issues/1",
                 "html_url": "https://github.com/random/random-repo/issues/1",
-                "title": "Found a bug", 
+                "title": "Found a bug",
                 "body": "I'm having a problem with this.",
                 "comments_url": "https://api.github.com/repos/random/random-repo/issues/1/comments",
-                "comments": 2
+                "comments": 2,
             }
         ]
         mock_get.return_value = mock_response
 
-        issues = self.github_manager._get_page_of_issues("https://api.github.com/repos/random/random-repo/issues?page=1").json()
+        issues = self.github_manager._get_page_of_issues(
+            "https://api.github.com/repos/random/random-repo/issues?page=1"
+        ).json()
 
-        assert len(issues) == 1  
+        assert len(issues) == 1
 
-    @patch('github.requests.get')
+    @patch("github.requests.get")
     def test_get_comments(self, mock_get):
         """Test retrieving comments for an issue."""
         mock_get.return_value.json.return_value = self.mock_comment_response().json()
-        
+
         comments = self.github_manager._get_comments("comments_url")
         assert len(comments) == 1
         assert comments[0].body == "This is a comment."
@@ -111,7 +132,7 @@ class TestGitHubIssuesManager:
             comments=[
                 GitHubIssueComment(url="comment_url_1", html_url="html_comment_url_1", body="First comment."),
                 GitHubIssueComment(url="comment_url_2", html_url="html_comment_url_2", body="Second comment."),
-            ]
+            ],
         )
 
         chunker = GitHubIssuesChunker(max_tokens=50)
@@ -124,7 +145,7 @@ class TestGitHubIssuesManager:
 
 
 class TestGitHubIssueComment:
-    
+
     def test_initialization(self):
         """Test the initialization of the GitHubIssueComment class."""
         comment = GitHubIssueComment(url="comment_url", html_url="html_comment_url", body="Sample comment")
@@ -139,10 +160,12 @@ class TestGitHubIssueComment:
 
 
 class TestGitHubIssue:
-    
+
     def test_initialization(self):
         """Test the initialization of the GitHubIssue class."""
-        issue = GitHubIssue(url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[])
+        issue = GitHubIssue(
+            url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[]
+        )
         assert issue.url == "issue_url"
         assert issue.html_url == "html_issue_url"
         assert issue.title == "Test Issue"
@@ -151,15 +174,19 @@ class TestGitHubIssue:
 
     def test_pretty_property(self):
         """Test the pretty property of the GitHubIssue class."""
-        issue = GitHubIssue(url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[])
+        issue = GitHubIssue(
+            url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[]
+        )
         assert issue.pretty == "# Issue: Test Issue\nTest Body"
 
 
 class TestIssueChunk:
-    
+
     def test_initialization(self):
         """Test the initialization of the IssueChunk class."""
-        issue = GitHubIssue(url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[])
+        issue = GitHubIssue(
+            url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[]
+        )
         chunk = IssueChunk(issue=issue, start_comment=0, end_comment=1)
         assert chunk.issue == issue
         assert chunk.start_comment == 0
@@ -167,32 +194,38 @@ class TestIssueChunk:
 
     def test_content_property(self):
         """Test the content property of the IssueChunk class."""
-        issue = GitHubIssue(url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[])
+        issue = GitHubIssue(
+            url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[]
+        )
         chunk = IssueChunk(issue=issue, start_comment=0, end_comment=1)
         assert chunk.content == "# Issue: Test Issue\nTest Body\n\n"
 
     def test_metadata_property(self):
         """Test the metadata property of the IssueChunk class."""
-        issue = GitHubIssue(url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[])
+        issue = GitHubIssue(
+            url="issue_url", html_url="html_issue_url", title="Test Issue", body="Test Body", comments=[]
+        )
         chunk = IssueChunk(issue=issue, start_comment=0, end_comment=1)
         expected_metadata = {
             "id": "html_issue_url_0_1",
             "url": "html_issue_url",
             "start_comment": 0,
             "end_comment": 1,
-            'text': '# Issue: Test Issue\nTest Body\n\n',
+            "text": "# Issue: Test Issue\nTest Body\n\n",
         }
         assert chunk.metadata == expected_metadata
 
     def test_num_tokens_property(self):
         """Test the num_tokens property of the IssueChunk class."""
-        issue = GitHubIssue(url="issue_url", html_url="html_issue_url", title="Test Issue", body="This is a test body.", comments=[])
+        issue = GitHubIssue(
+            url="issue_url", html_url="html_issue_url", title="Test Issue", body="This is a test body.", comments=[]
+        )
         chunk = IssueChunk(issue=issue, start_comment=0, end_comment=1)
-        assert chunk.num_tokens == 12 
+        assert chunk.num_tokens == 12
 
 
 class TestGitHubIssuesChunker:
-    
+
     def test_initialization(self):
         """Test the initialization of the GitHubIssuesChunker class."""
         chunker = GitHubIssuesChunker(max_tokens=50)
@@ -205,16 +238,16 @@ class TestGitHubIssuesChunker:
             html_url="html_issue_url",
             title="Test Issue",
             body="This is a long body of the issue that needs to be chunked.",
-            comments=[]
+            comments=[],
         )
-        
+
         chunker = GitHubIssuesChunker(max_tokens=50)
         chunks = chunker.chunk(content=issue, metadata={})
-        
+
         assert len(chunks) > 0
 
         assert all(isinstance(chunk, IssueChunk) for chunk in chunks)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     pytest.main()