Code Indexer Loop is a Python library for indexing and retrieving source code files through an integrated vector database that's continuously and efficiently updated.
import atexit
import os
from pathlib import Path

import chromadb
from langchain.embeddings.openai import OpenAIEmbeddings
from llama_index import ServiceContext, VectorStoreIndex
from llama_index.embeddings import LangchainEmbedding
from llama_index.schema import NodeWithScore, TextNode
from llama_index.vector_stores import ChromaVectorStore
from watchdog.events import FileSystemEventHandler
from watchdog.observers import Observer

from code_indexer_loop.code_splitter import CodeSplitter
from code_indexer_loop.constants import EXTENSION_TO_TREE_SITTER_LANGUAGE
from code_indexer_loop.utils import hash_md5


class CodeIndexer:
    src_dir: str
    target_chunk_tokens: int
    max_chunk_tokens: int
    enforce_max_chunk_tokens: bool
    token_model: str
    code_splitters = {}
    hash_cache = {}
    index: VectorStoreIndex = None

    def __init__(
        self,
        src_dir: str,
        target_chunk_tokens: int = 300,
        max_chunk_tokens: int = 1000,
        enforce_max_chunk_tokens: bool = False,
        coalesce: int = 50,
        token_model: str = "gpt-4",
        watch: bool = False,
    ):
        self.src_dir = src_dir
        self.target_chunk_tokens = target_chunk_tokens
        self.max_chunk_tokens = max_chunk_tokens
        self.enforce_max_chunk_tokens = enforce_max_chunk_tokens
        self.coalesce = coalesce
        self.token_model = token_model
        self._create_index()
        self.refresh_nodes()

        if watch:
            self._start_watching()
            atexit.register(self._stop_watching)

    def query(self, query: str, k=10) -> str:
        return "\n".join(
            [node_with_score.node.text for node_with_score in self.index.as_retriever(similarity_top_k=k).retrieve(query)]
        )

    def query_nodes(self, query: str, k=10) -> list[NodeWithScore]:
        return self.index.as_retriever(similarity_top_k=k).retrieve(query)

    def query_documents(self, query: str, k=10) -> list[dict[str, str]]:
        nodes = self.index.as_retriever(similarity_top_k=k).retrieve(query)
        files = [node_with_score.node.metadata["file"] for node_with_score in nodes]
        # Deduplicate files, preserving order
        files = list(dict.fromkeys(files))
        # Read file contents
        contents = []
        for file in files:
            with open(file, "r") as f:
                contents.append(
                    {
                        "file": file,
                        "content": f.read(),
                    }
                )
        return contents

    def add_file(self, file: str):
        ext = os.path.splitext(file)[1]
        text_splitter = self._get_code_splitter(ext)

        calculated_hash = hash_md5(file)
        if file in self.hash_cache:
            if self.hash_cache[file] == calculated_hash:
                # Skip file if it hasn't changed
                return
        else:
            self.hash_cache[file] = calculated_hash

        with open(file, "r") as f:
            text = f.read()
            nodes = [
                TextNode(
                    text=chunk,
                    metadata={
                        "file": file,
                    },
                )
                for chunk in text_splitter.split_text(text)
            ]

            self._remove_old_nodes(file)
            self._insert_nodes(nodes)

    def remove_file(self, file: str):
        self._remove_old_nodes(file)
        del self.hash_cache[file]

    def refresh_nodes(self):
        files = self._find_files(self.src_dir, EXTENSION_TO_TREE_SITTER_LANGUAGE)

        # Clear any files that no longer exist
        for file in list(self.hash_cache.keys()):
            if file not in files:
                del self.hash_cache[file]
                self._remove_old_nodes(file)

        # For each file, split into chunks and index
        for file in files:
            self.add_file(str(file))

    def _start_watching(self):
        event_handler = CodeChangeHandler(self)
        self.observer = Observer()
        self.observer.schedule(event_handler, self.src_dir, recursive=True)
        self.observer.start()

    def _stop_watching(self):
        if hasattr(self, "observer"):
            self.observer.stop()
            self.observer.join()

    def _find_files(self, path, include_ext={}):
        """
        Recursively find all files in a given path.

        Parameters:
            path (str): The root directory to start searching from.
            include_ext (dict): A dictionary of file extensions to include
                (keys are extensions including leading period if applicable).

        Returns:
            list: A list of full file paths for each file found.
        """
        found_files = []

        for root, _, files in os.walk(path):
            for file in files:
                # Check if the file should be excluded based on its extension
                file_ext = os.path.splitext(file)[1]
                if file_ext in include_ext:
                    # Construct the full path of the file and append to list
                    full_path = Path(os.path.join(root, file)).resolve()
                    found_files.append(full_path)

        return set(found_files)

    def _get_code_splitter(self, ext) -> CodeSplitter:
        if ext not in EXTENSION_TO_TREE_SITTER_LANGUAGE:
            raise ValueError(f"Extension {ext} not supported.")
        language = EXTENSION_TO_TREE_SITTER_LANGUAGE[ext]
        if language not in self.code_splitters:
            text_splitter = CodeSplitter(
                language=language,
                target_chunk_tokens=self.target_chunk_tokens,
                max_chunk_tokens=self.max_chunk_tokens,
                enforce_max_chunk_tokens=self.enforce_max_chunk_tokens,
                coalesce=self.coalesce,
                token_model=self.token_model,
            )
            self.code_splitters[ext] = text_splitter

        return self.code_splitters[ext]

    def _remove_old_nodes(self, file):
        # Remove existing nodes for the same file
        self.index.vector_store.client.delete(where={"file": file})

    def _insert_nodes(self, nodes):
        self.index.insert_nodes(nodes)

    def _create_index(self) -> VectorStoreIndex:
        # Create client and a new collection
        chroma_client = chromadb.EphemeralClient()
        chroma_collection = chroma_client.create_collection("code-index")

        # Define embedding function
        embed_model = LangchainEmbedding(OpenAIEmbeddings())

        # Set up ChromaVectorStore and load in data
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        service_context = ServiceContext.from_defaults(embed_model=embed_model)
        index = VectorStoreIndex.from_vector_store(vector_store=vector_store, service_context=service_context)

        self.index = index
        return index


class CodeChangeHandler(FileSystemEventHandler):
    def __init__(self, indexer: CodeIndexer):
        self.indexer = indexer

    def on_modified(self, event):
        if event.is_directory:
            # Directory modifications shouldn't trigger a reindex
            return
        else:
            # Update only if the changed file has a supported extension
            ext = os.path.splitext(event.src_path)[1]
            if ext in EXTENSION_TO_TREE_SITTER_LANGUAGE:
                self.indexer.add_file(event.src_path)

    def on_created(self, event):
        if event.is_directory:
            self.indexer.refresh_nodes()
        else:
            # Update only if the changed file has a supported extension
            ext = os.path.splitext(event.src_path)[1]
            if ext in EXTENSION_TO_TREE_SITTER_LANGUAGE:
                self.indexer.add_file(event.src_path)

    def on_moved(self, event):
        self.indexer.refresh_nodes()

    def on_deleted(self, event):
        if event.is_directory:
            self.indexer.refresh_nodes()
        else:
            ext = os.path.splitext(event.src_path)[1]
            if ext in EXTENSION_TO_TREE_SITTER_LANGUAGE:
                self.indexer.remove_file(event.src_path)

code_indexer_loop/api.py
definitive-io/code-indexer-loop
"""Code Splitter.

Implementation amalgamated from:
https://docs.sweep.dev/blogs/chunking-improvements
https://docs.sweep.dev/blogs/chunking-2m-files
https://github.com/jerryjliu/llama_index/pull/7100

"""

import re
from dataclasses import dataclass
from typing import List, Optional, Union

import tiktoken
from tree_sitter import Node


class MaxChunkLengthExceededError(Exception):
    pass


@dataclass
class Span:
    # Represents a slice of a string
    start: int = 0
    end: int = 0

    def __post_init__(self):
        # If end is None, set it to start
        if self.end is None:
            self.end = self.start

    def extract(self, s: bytes) -> bytes:
        # Grab the corresponding substring of string s by bytes
        return s[self.start : self.end]

    def extract_lines(self, s: str) -> str:
        lines = s.split("\n")
        selected_lines = lines[self.start : self.end]
        joined = "\n".join(selected_lines)
        # if selection doesn't extend to the last line, add the missing newline
        if self.end < len(lines):
            joined += "\n"
        return joined

    def __add__(self, other: Union["Span", int]) -> "Span":
        # e.g. Span(1, 2) + Span(2, 4) = Span(1, 4) (concatenation)
        # There are no safety checks: Span(a, b) + Span(c, d) = Span(a, d)
        # and there are no requirements for b = c.
        if isinstance(other, int):
            return Span(self.start + other, self.end + other)
        elif isinstance(other, Span):
            return Span(self.start, other.end)
        else:
            raise NotImplementedError()

    def __len__(self) -> int:
        # i.e. Span(a, b) = b - a
        return self.end - self.start


class TokenCounter:
    default_model: str
    initialized_models = {}

    def __init__(self, default_model: str):
        self.default_model = default_model

    def count(self, text: str, model: Optional[str] = None):
        if model is None:
            model = self.default_model

        if model not in self.initialized_models:
            try:
                self.initialized_models[model] = tiktoken.encoding_for_model(model)
            except KeyError:
                raise KeyError(f"Model {model} not supported.")

        return len(self.initialized_models[model].encode(text, disallowed_special=()))

    def count_chunk(self, chunk: Span, source_code: bytes, model: Optional[str] = None):
        return self.count(chunk.extract(source_code).decode("utf-8"), model)


class CodeSplitter:
    """Split code using a AST parser."""

    language: str
    target_chunk_tokens: int
    max_chunk_tokens: int
    enforce_max_chunk_tokens: bool
    coalesce: int
    token_counter: TokenCounter

    def __init__(
        self,
        language: str,
        target_chunk_tokens: int,
        max_chunk_tokens: int,
        enforce_max_chunk_tokens: bool,
        coalesce: int,
        token_model: str,
    ):
        self.token_counter = TokenCounter(default_model=token_model)
        self.target_chunk_tokens = target_chunk_tokens
        self.max_chunk_tokens = max_chunk_tokens
        self.enforce_max_chunk_tokens = enforce_max_chunk_tokens
        self.language = language
        self.coalesce = coalesce

    @classmethod
    def class_name(cls) -> str:
        """Get class name."""
        return "CodeSplitter"

    def chunk_tree(
        self,
        tree,
        source_code: bytes,
    ) -> list[Span]:
        # 1. Recursively form chunks
        def chunk_node(node: Node) -> list[Span]:
            chunks: list[Span] = []
            current_chunk: Span = Span(node.start_byte, node.start_byte)
            node_children = node.children
            for child in node_children:
                child_token_len = self.token_counter.count_chunk(Span(child.start_byte, child.end_byte), source_code)
                child_and_current_token_len = self.token_counter.count_chunk(
                    Span(child.start_byte, child.end_byte), source_code
                ) + self.token_counter.count_chunk(current_chunk, source_code)

                if child_token_len > self.target_chunk_tokens:
                    if child_token_len > self.max_chunk_tokens and self.enforce_max_chunk_tokens:
                        raise MaxChunkLengthExceededError(
                            f"Chunk token length {child_token_len} exceeds maximum {self.max_chunk_tokens}."
                        )

                    chunks.append(current_chunk)
                    current_chunk = Span(child.end_byte, child.end_byte)
                    chunks.extend(chunk_node(child))
                elif child_and_current_token_len > self.target_chunk_tokens:
                    if child_and_current_token_len > self.max_chunk_tokens and self.enforce_max_chunk_tokens:
                        raise MaxChunkLengthExceededError(
                            f"Chunk token length {child_and_current_token_len}"
                            f" exceeds maximum {self.max_chunk_tokens}."
                        )
                    chunks.append(current_chunk)
                    current_chunk = Span(child.start_byte, child.end_byte)
                else:
                    current_chunk += Span(child.start_byte, child.end_byte)

            final_chunk_token_len = self.token_counter.count_chunk(current_chunk, source_code)
            if final_chunk_token_len > self.max_chunk_tokens and self.enforce_max_chunk_tokens:
                raise MaxChunkLengthExceededError(
                    f"Chunk token length {final_chunk_token_len} exceeds maximum {self.max_chunk_tokens}."
                )
            chunks.append(current_chunk)
            return chunks

        chunks = chunk_node(tree.root_node)

        # Filter empty chunks
        chunks = [chunk for chunk in chunks if len(chunk) > 0]

        # Early return if there is no chunk
        if len(chunks) == 0:
            return []
        # Early return if there is only one chunk
        if len(chunks) < 2:
            return [Span(0, len(chunks[0]))]

        # Filling in the gaps
        # by aligning end of one chunk with start of next
        chunks[0].start = 0
        for prev, curr in zip(chunks[:-1], chunks[1:]):
            prev.end = curr.start
        curr.end = len(source_code)

        # Combining small chunks with bigger ones
        new_chunks = []
        aggregated_chunk = Span(0, 0)
        aggregated_chunk_token_len = 0
        for chunk in chunks:
            # Check if the combined chunk exceeds target_chunk_tokens
            # Note, at this point no chunk exceeds max_chunk_tokens
            # if max_chunk_tokens is enforced.
            chunk_token_len = self.token_counter.count_chunk(chunk, source_code)
            if chunk_token_len > self.target_chunk_tokens:
                new_chunks.append(aggregated_chunk)
                new_chunks.append(chunk)
                aggregated_chunk = Span(chunk.end, chunk.end)
                aggregated_chunk_token_len = 0
            elif aggregated_chunk_token_len + chunk_token_len > self.target_chunk_tokens:
                new_chunks.append(aggregated_chunk)
                aggregated_chunk = Span(chunk.start, chunk.end)
                aggregated_chunk_token_len = chunk_token_len
            else:
                # Combined chunk does not exceed target_chunk_tokens
                # so we add the current chunk to the aggregated_chunk.
                # Note, there is no need to check whether the combined chunk
                # exceeds max_chunk_tokens because we have already checked.
                aggregated_chunk += chunk
                aggregated_chunk_token_len += chunk_token_len
                if aggregated_chunk_token_len > self.coalesce:
                    new_chunks.append(aggregated_chunk)
                    aggregated_chunk = Span(chunk.end, chunk.end)
                    aggregated_chunk_token_len = 0

        if len(aggregated_chunk) > 0:
            new_chunks.append(aggregated_chunk)

        # Changing line numbers
        line_chunks = [
            Span(
                self.get_line_number(chunk.start, source_code),
                self.get_line_number(chunk.end, source_code),
            )
            for chunk in new_chunks
        ]

        # Eliminating empty chunks
        line_chunks = [chunk for chunk in line_chunks if len(chunk) > 0]
        return line_chunks

    def split_and_keep_newline(self, byte_str):
        return re.split(b"(?<=\n)", byte_str)

    def get_line_number(self, index: int, source_code: bytes) -> int:
        total_chars = 0
        for line_number, line in enumerate(self.split_and_keep_newline(source_code), start=1):
            total_chars += len(line)
            if total_chars > index:
                return line_number - 1
        return line_number

    def split_text(self, text: str) -> List[str]:
        """Split incoming code and return chunks using the AST."""
        try:
            import tree_sitter_languages
        except ImportError:
            raise ImportError("Please install tree_sitter_languages to use CodeSplitter.")

        try:
            parser = tree_sitter_languages.get_parser(self.language)
        except Exception as e:
            print(
                f"Could not get parser for language {self.language}. Check "
                "https://github.com/grantjenks/py-tree-sitter-languages#license "
                "for a list of valid languages."
            )
            raise e

        tree = parser.parse(text.encode("utf-8"))
        if not tree.root_node.children or tree.root_node.children[0].type != "ERROR":
            line_spans = self.chunk_tree(tree, text.encode("utf-8"))
            chunks = [line_span.extract_lines(text) for line_span in line_spans]
            return chunks
        else:
            raise ValueError(f"Could not parse code with language {self.language}.")

code_indexer_loop/code_splitter.py
definitive-io/code-indexer-loop
EXTENSION_TO_TREE_SITTER_LANGUAGE = {
    ".c": "c",
    ".cc": "cpp",
    ".cpp": "cpp",
    ".cs": "c-sharp",
    ".cxx": "cpp",
    ".go": "go",
    ".hs": "haskell",
    ".java": "java",
    ".jl": "julia",
    ".js": "javascript",
    ".jsx": "javascript",
    ".php": "php",
    ".py": "python",
    ".rb": "ruby",
    ".rs": "rust",
    ".scala": "scala",
    ".sql": "sql",
    ".swift": "swift",
    ".ts": "typescript",
    ".tsx": "typescript",
}

code_indexer_loop/constants.py
definitive-io/code-indexer-loop
import os

import pytest

from code_indexer_loop.code_splitter import (CodeSplitter,
                                             MaxChunkLengthExceededError,
                                             TokenCounter)

THIS_FILE_DIR = os.path.dirname(os.path.realpath(__file__))


def create_code_splitter(language="python", target_chunk_tokens=5, max_chunk_tokens=200, enforce_max_chunk_tokens=True):
    return CodeSplitter(
        language=language,
        target_chunk_tokens=target_chunk_tokens,
        max_chunk_tokens=max_chunk_tokens,
        enforce_max_chunk_tokens=enforce_max_chunk_tokens,
        token_model="gpt-4",
        coalesce=50,
    )


def test_code_splitter_prefix_model():
    CodeSplitter(
        language="python",
        target_chunk_tokens=10,
        max_chunk_tokens=10,
        enforce_max_chunk_tokens=True,
        token_model="gpt-4-32k-0613",
        coalesce=50,
    )


def test_code_splitter():
    python_code_splitter = create_code_splitter()
    chunks = python_code_splitter.split_text(
        """def foo():
    print("Hello, world!")

print(1)"""
    )
    assert chunks[0].startswith("def foo():")
    assert not chunks[0].endswith('")')


def test_code_splitter_newlines():
    python_code_splitter = create_code_splitter()
    chunks = python_code_splitter.split_text(
        """
def foo():
    print("Hello, world!")

print(1)

"""
    )
    assert chunks[0].startswith("\ndef foo():")
    assert not chunks[0].endswith('")')
    assert chunks[-1].endswith("\n\n")


def test_code_splitter_raise():
    python_code_splitter = create_code_splitter(max_chunk_tokens=5)
    with pytest.raises(MaxChunkLengthExceededError):
        python_code_splitter.split_text(
            """
def mostdefinitelynotlessthan5tokens():
    pass
"""
        )


def test_code_splitter_noraise():
    python_code_splitter = create_code_splitter(max_chunk_tokens=5, enforce_max_chunk_tokens=False)
    python_code_splitter.split_text(
        """
def mostdefinitelynotlessthan5tokens():
    pass
"""
    )


def test_code_splitter_token_lengths():
    tc = TokenCounter(default_model="gpt-4")
    max_chunk_tokens = 20
    python_code_splitter = create_code_splitter(
        max_chunk_tokens=max_chunk_tokens, target_chunk_tokens=max_chunk_tokens // 2
    )
    source_code = """
def add(a, b):
    return a + b

def subtract(a, b):
    return a - b

add(1, 2)
"""
    chunks = python_code_splitter.split_text(source_code)
    joined_chunks = "".join(chunks)
    assert source_code == joined_chunks

    chunk_lengths = [tc.count(chunk) for chunk in chunks]
    assert all([chunk_length <= max_chunk_tokens for chunk_length in chunk_lengths])


def test_long_file():
    hard_file_path = os.path.join(THIS_FILE_DIR, "test_api_dummy_file.py.txt")
    with open(hard_file_path, "r") as f:
        source_code = f.read()

    python_code_splitter = create_code_splitter(target_chunk_tokens=1000, max_chunk_tokens=9000)
    chunks = python_code_splitter.split_text(source_code)
    joined_chunks = "".join(chunks)
    assert source_code == joined_chunks


def test_sql():
    sql_file_path = os.path.join(THIS_FILE_DIR, "test_api_dummy_sql.sql.txt")
    with open(sql_file_path, "r") as f:
        source_code = f.read()

    sql_code_splitter = CodeSplitter(
        language="sql",
        target_chunk_tokens=10,
        max_chunk_tokens=1000,
        enforce_max_chunk_tokens=True,
        token_model="gpt-4",
        coalesce=50,
    )

    chunks = sql_code_splitter.split_text(source_code)
    joined_chunks = "".join(chunks)
    assert source_code == joined_chunks

code_indexer_loop/test_api.py
definitive-io/code-indexer-loop
import hashlib


def hash_md5(filename):
    hash_md5 = hashlib.md5()
    with open(filename, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()

code_indexer_loop/utils.py
definitive-io/code-indexer-loop