Local first semantic code search and chat powered by vector embeddings and LLMs
from codeqai import app


def main():
    app.run()


if __name__ == "__main__":
    main()

codeqai/__main__.py
fynnfluegge/codeqai
import argparse
import os

from dotenv import dotenv_values, load_dotenv
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationSummaryMemory
from rich.console import Console
from rich.markdown import Markdown
from rich.syntax import Syntax
from yaspin import yaspin

from codeqai import codeparser, repo, utils
from codeqai.config import (
    create_cache_dir,
    create_config,
    get_cache_path,
    get_config_path,
    load_config,
)
from codeqai.constants import EmbeddingsModel, LlmHost
from codeqai.embeddings import Embeddings
from codeqai.llm import LLM
from codeqai.vector_store import VectorStore


def env_loader(env_path, required_keys=None):
    """
    Args :
    env_path = source path of .env file.
    required_keys = ["OPENAI_KEY"] #change this according to need

    #running/calling the function.
    configs = env_loader('.env', required_keys)
    """

    # create env file if does not exists
    # parse required keys in the file if it's not None
    if not os.path.exists(env_path) or os.path.getsize(env_path) == 0:
        with open(env_path, "w") as env_f:
            if required_keys:
                for key in required_keys:
                    env_f.write(f'{key}=""\n')
            else:
                pass

    configs = dotenv_values(env_path)
    changed = False
    for key, value in configs.items():
        env_key = os.getenv(key)
        if not value and not env_key:
            value = input(
                f"[+] Key {utils.get_bold_text(key)} is required. Please enter it's value: "
            )
            configs[key] = value
            changed = True
        elif not value and env_key:
            value = env_key
            configs[key] = value
            changed = True

    # update the .env file if config is changed/taken from user
    if changed:
        with open(env_path, "w") as env_f:
            for key, value in configs.items():
                env_f.write(f'{key}="{value}"\n')

    load_dotenv(env_path, override=True)


def run():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "action",
        choices=["search", "chat", "configure"],
        help="Action to perform. 'search' will semantically search the codebase. 'chat' will chat with the codebase.",
    )
    args = parser.parse_args()

    if args.action == "configure":
        create_config()
        exit()

    # load config
    config = {}
    try:
        config = load_config()
    except FileNotFoundError:
        config = create_config()

    # lookup env variables
    required_keys = []
    if (
        config["llm-host"] == LlmHost.OPENAI.value
        or config["embeddings"] == EmbeddingsModel.OPENAI_TEXT_EMBEDDING_ADA_002.value
    ):
        required_keys.append("OPENAI_API_KEY")

    if (
        config["llm-host"] == LlmHost.AZURE_OPENAI.value
        or config["embeddings"] == EmbeddingsModel.AZURE_OPENAI.value
    ):
        required_keys.extend(
            [
                "OPENAI_API_TYPE",
                "OPENAI_API_BASE_URL",
                "OPENAI_API_KEY",
                "OPENAI_API_VERSION",
            ]
        )
    env_path = get_config_path().replace("config.yaml", ".env")
    env_loader(env_path, required_keys)

    repo_name = repo.get_git_root(os.getcwd()).split("/")[-1]

    # init cache
    create_cache_dir()

    embeddings_model = Embeddings(
        local=config["local"],
        model=EmbeddingsModel[config["embeddings"].upper().replace("-", "_")],
        deployment=config["embeddings-deployment"]
        if "embeddings-deployment" in config
        else None,
    )

    # check if faiss.index exists
    if not os.path.exists(os.path.join(get_cache_path(), f"{repo_name}.faiss")):
        # sync repo
        spinner = yaspin(text="🔧 Parsing codebase...", color="green")
        files = repo.load_files()
        documents = codeparser.parse_code_files(files)
        spinner.stop()
        vector_store = VectorStore(
            repo_name,
            embeddings=embeddings_model.embeddings,
            documents=documents,
        )
    else:
        vector_store = VectorStore(repo_name, embeddings=embeddings_model.embeddings)

    llm = LLM(
        llm_host=LlmHost[config["llm-host"].upper().replace("-", "_")],
        chat_model=config["chat-model"],
        deployment=config["model-deployment"] if "model-deployment" in config else None,
    )
    memory = ConversationSummaryMemory(
        llm=llm.chat_model, memory_key="chat_history", return_messages=True
    )
    qa = ConversationalRetrievalChain.from_llm(
        llm.chat_model, retriever=vector_store.retriever, memory=memory
    )

    console = Console()
    while True:
        choice = None
        if args.action == "search":
            search_pattern = input("🔎 Enter a search pattern: ")
            spinner = yaspin(text="🤖 Processing...", color="green")
            spinner.start()
            similarity_result = vector_store.similarity_search(search_pattern)
            spinner.stop()
            for doc in similarity_result:
                language = utils.get_programming_language(
                    utils.get_file_extension(doc.metadata["filename"])
                )

                syntax = Syntax(
                    doc.page_content,
                    language.value,
                    theme="monokai",
                    line_numbers=True,
                )
                print(doc.metadata["filename"] + " -> " + doc.metadata["method_name"])
                console.print(syntax)
                print()

            choice = input("[?] (C)ontinue search or (E)xit [C]:").strip().lower()

        elif args.action == "chat":
            question = input("💬 Ask anything about the codebase: ")
            spinner = yaspin(text="🤖 Processing...", color="green")
            spinner.start()
            result = qa(question)
            spinner.stop()
            markdown = Markdown(result["answer"])
            console.print(markdown)

            choice = (
                input("[?] (C)ontinue chat, (R)eset chat or (E)xit [C]:")
                .strip()
                .lower()
            )

            if choice == "r":
                memory.clear()
                print("Chat history cleared.")

        if choice == "" or choice == "c":
            continue
        elif choice == "e":
            break
        else:
            print("Invalid choice. Please enter 'C', 'E', or 'R'.")

codeqai/app.py
fynnfluegge/codeqai
import os

from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter

from codeqai import utils
from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter, TreesitterMethodNode


def parse_code_files(code_files: list[str]) -> list[Document]:
    documents = []
    code_splitter = None
    for code_file in code_files:
        with open(code_file, "r") as file:
            file_bytes = file.read().encode()

            file_extension = utils.get_file_extension(code_file)
            programming_language = utils.get_programming_language(file_extension)
            if programming_language == Language.UNKNOWN:
                continue

            langchain_language = utils.get_langchain_language(programming_language)

            if langchain_language:
                code_splitter = RecursiveCharacterTextSplitter.from_language(
                    language=langchain_language,
                    chunk_size=512,
                    chunk_overlap=128,
                )

            treesitter_parser = Treesitter.create_treesitter(programming_language)
            treesitterNodes: list[TreesitterMethodNode] = treesitter_parser.parse(
                file_bytes
            )
            for node in treesitterNodes:
                method_source_code = node.method_source_code
                filename = os.path.basename(code_file)

                if node.doc_comment and programming_language != Language.PYTHON:
                    method_source_code = node.doc_comment + "\n" + method_source_code

                splitted_documents = [method_source_code]
                if code_splitter:
                    splitted_documents = code_splitter.split_text(method_source_code)

                for splitted_document in splitted_documents:
                    document = Document(
                        page_content=splitted_document,
                        metadata={
                            "filename": filename,
                            "method_name": node.name,
                        },
                    )
                    documents.append(document)

    return documents

codeqai/codeparser.py
fynnfluegge/codeqai
import os
import platform
from pathlib import Path

import inquirer
import yaml

from codeqai.constants import EmbeddingsModel, LlmHost


def get_cache_path():
    system = platform.system()

    if system == "Linux" or system == "Darwin":
        user_home = os.path.expanduser("~")
        cache_dir = os.path.join(user_home, ".cache", "codeqai")
    elif system == "Windows":
        user_home = os.path.expanduser("~")
        cache_dir = os.path.join(user_home, "AppData", "Local", "codeqai")
    else:
        raise NotImplementedError(f"Unsupported platform: {system}")

    return cache_dir


def create_cache_dir():
    if not os.path.exists(get_cache_path()):
        path = Path(get_cache_path())
        path.mkdir(parents=True, exist_ok=True)


def get_config_path():
    system = platform.system()

    if system == "Linux" or system == "Darwin":
        user_home = os.path.expanduser("~")
        config_dir = os.path.join(user_home, ".config", "codeqai")
    elif system == "Windows":
        user_home = os.path.expanduser("~")
        config_dir = os.path.join(user_home, "AppData", "Roaming", "codeqai")
    else:
        raise NotImplementedError(f"Unsupported platform: {system}")

    config_file_path = os.path.join(config_dir, "config.yaml")

    return config_file_path


def load_config():
    with open(get_config_path(), "r") as config_file:
        config = yaml.safe_load(config_file)
    return config


def save_config(config):
    with open(get_config_path(), "w") as config_file:
        yaml.dump(config, config_file, default_flow_style=False)


def create_config():
    os.makedirs(os.path.dirname(get_config_path()), exist_ok=True)

    questions = [
        inquirer.Confirm(
            "confirm", message="Do you want to use local models?", default=False
        ),
    ]

    confirm = inquirer.prompt(questions)

    if confirm and confirm["confirm"]:
        questions = [
            inquirer.List(
                "embeddings",
                message="Which local embeddings model do you want to use?",
                choices=[
                    EmbeddingsModel.INSTRUCTOR_LARGE.value,
                    EmbeddingsModel.SENTENCETRANSFORMERS_ALL_MPNET_BASE_V2.value,
                    EmbeddingsModel.SENTENCETRANSFORMERS_ALL_MINILM_L6_V2.value,
                ],
                default=EmbeddingsModel.INSTRUCTOR_LARGE.value,
            ),
            inquirer.List(
                "llm-host",
                message="Which local LLM host do you want to use?",
                choices=[
                    LlmHost.LLAMACPP.value,
                    LlmHost.OLLAMA.value,
                ],
                default=LlmHost.LLAMACPP.value,
            ),
        ]
    else:
        questions = [
            inquirer.List(
                "embeddings",
                message="Which remote embeddings do you want to use?",
                choices=[
                    EmbeddingsModel.OPENAI_TEXT_EMBEDDING_ADA_002.value,
                    EmbeddingsModel.AZURE_OPENAI.value,
                ],
                default=EmbeddingsModel.OPENAI_TEXT_EMBEDDING_ADA_002.value,
            ),
            inquirer.List(
                "llm-host",
                message="Which remote LLM do you want to use?",
                choices=[
                    LlmHost.OPENAI.value,
                    LlmHost.AZURE_OPENAI.value,
                ],
                default=LlmHost.OPENAI.value,
            ),
        ]

    answers = inquirer.prompt(questions)

    if confirm and answers:
        config = {
            "local": confirm["confirm"],
            "embeddings": answers["embeddings"],
            "llm-host": answers["llm-host"],
        }

        if answers["embeddings"] == EmbeddingsModel.AZURE_OPENAI.value:
            questions = [
                inquirer.Text(
                    "deployment",
                    message="Please enter the Azure OpenAI embeddings deployment name.",
                    default="",
                ),
            ]
            deployment_answer = inquirer.prompt(questions)
            if deployment_answer and deployment_answer["deployment"]:
                config["embeddings-deployment"] = deployment_answer["deployment"]

        if answers["llm-host"] == LlmHost.AZURE_OPENAI.value:
            questions = [
                inquirer.Text(
                    "deployment",
                    message="Please enter the Azure OpenAI model deployment name",
                    default="",
                ),
            ]
            deployment_answer = inquirer.prompt(questions)
            if deployment_answer and deployment_answer["deployment"]:
                config["model-deployment"] = deployment_answer["deployment"]

        if answers["llm-host"] == LlmHost.LLAMACPP.value:
            questions = [
                inquirer.Text(
                    "chat-model",
                    message="Please enter the path to the LLM model",
                    default="",
                ),
            ]
        elif answers["llm-host"] == LlmHost.OLLAMA.value:
            questions = [
                inquirer.List(
                    "chat-model",
                    message="Which Ollama chat model do you want to use?",
                    choices=[
                        "llama2",
                        "llama2:13b",
                        "llama2:70b",
                        "codellama",
                    ],
                    default="llama2:13b",
                ),
            ]
        elif answers["llm-host"] == "OpenAI":
            questions = [
                inquirer.List(
                    "chat-model",
                    message="Which OpenAI chat model do you want to use?",
                    choices=[
                        "gpt-3.5-turbo",
                        "gpt-3.5-turbo-16k",
                        "gpt-4",
                    ],
                    default="gpt-3.5-turbo",
                ),
            ]

        answers = inquirer.prompt(questions)
        if answers and answers["chat-model"]:
            config["chat-model"] = answers["chat-model"]

        save_config(config)

        return config

    return {}

codeqai/config.py
fynnfluegge/codeqai
from enum import Enum


class Language(Enum):
    PYTHON = "python"
    JAVASCRIPT = "javascript"
    TYPESCRIPT = "typescript"
    JAVA = "java"
    CPP = "cpp"
    C = "c"
    GO = "go"
    RUST = "rust"
    KOTLIN = "kotlin"
    C_SHARP = "c_sharp"
    OBJECTIVE_C = "objective_c"
    SCALA = "scala"
    LUA = "lua"
    UNKNOWN = "unknown"


class EmbeddingsModel(Enum):
    SENTENCETRANSFORMERS_ALL_MPNET_BASE_V2 = "SentenceTransformers-all-mpnet-base-v2"
    SENTENCETRANSFORMERS_ALL_MINILM_L6_V2 = "SentenceTransformers-all-MiniLM-L6-v2"
    INSTRUCTOR_LARGE = "Instructor-Large"
    OPENAI_TEXT_EMBEDDING_ADA_002 = "OpenAI-text-embedding-ada-002"
    AZURE_OPENAI = "Azure-OpenAI"


class LlmHost(Enum):
    LLAMACPP = "Llamacpp"
    OLLAMA = "Ollama"
    OPENAI = "OpenAI"
    AZURE_OPENAI = "Azure-OpenAI"

codeqai/constants.py
fynnfluegge/codeqai
import inquirer
from langchain.embeddings import (HuggingFaceEmbeddings,
                                  HuggingFaceInstructEmbeddings)
from langchain.embeddings.openai import OpenAIEmbeddings

from codeqai import utils
from codeqai.constants import EmbeddingsModel


class Embeddings:
    def __init__(
        self,
        local=False,
        model=EmbeddingsModel.OPENAI_TEXT_EMBEDDING_ADA_002,
        deployment=None,
    ):
        if not local:
            if model == EmbeddingsModel.OPENAI_TEXT_EMBEDDING_ADA_002:
                self.embeddings = OpenAIEmbeddings(
                    client=None, model="text-embedding-ada-002"
                )
            elif model == EmbeddingsModel.AZURE_OPENAI and deployment:
                self.embeddings = OpenAIEmbeddings(client=None, deployment=deployment)
        else:
            try:
                import sentence_transformers  # noqa: F401
            except ImportError:
                self._install_sentence_transformers()

            if model == EmbeddingsModel.SENTENCETRANSFORMERS_ALL_MPNET_BASE_V2:
                self.embeddings = HuggingFaceEmbeddings()
            elif model == EmbeddingsModel.SENTENCETRANSFORMERS_ALL_MINILM_L6_V2:
                self.embeddings = HuggingFaceEmbeddings(
                    model_name=EmbeddingsModel.SENTENCETRANSFORMERS_ALL_MINILM_L6_V2.value.replace(
                        "SentenceTransformers-", ""
                    )
                )
            elif model == EmbeddingsModel.INSTRUCTOR_LARGE:
                try:
                    from InstructorEmbedding import INSTRUCTOR  # noqa: F401
                except ImportError:
                    self._install_instructor_embedding()

                self.embeddings = HuggingFaceInstructEmbeddings()

    def _install_sentence_transformers(self):
        question = [
            inquirer.Confirm(
                "confirm",
                message=f"{utils.get_bold_text('SentenceTransformers')} not found in this python environment. Do you want to install it now?",
                default=True,
            ),
        ]

        answers = inquirer.prompt(question)
        if answers and answers["confirm"]:
            import subprocess
            import sys

            try:
                subprocess.run(
                    [
                        sys.executable,
                        "-m",
                        "pip",
                        "install",
                        "sentence_transformers",
                    ],
                    check=True,
                )
            except subprocess.CalledProcessError as e:
                print(f"Error during sentence_transformers installation: {e}")
        else:
            exit("sentence_transformers is required for local embeddings.")

    def _install_instructor_embedding(self):
        question = [
            inquirer.Confirm(
                "confirm",
                message=f"{utils.get_bold_text('InstructorEmbedding')} not found in this python environment. Do you want to install it now?",
                default=True,
            ),
        ]

        answers = inquirer.prompt(question)
        if answers and answers["confirm"]:
            import subprocess
            import sys

            try:
                subprocess.run(
                    [
                        sys.executable,
                        "-m",
                        "pip",
                        "install",
                        "InstructorEmbedding",
                    ],
                    check=True,
                )
            except subprocess.CalledProcessError as e:
                print(f"Error during sentence_transformers installation: {e}")
        else:
            exit("InstructorEmbedding is required for local embeddings.")

codeqai/embeddings.py
fynnfluegge/codeqai
import os
import subprocess
import sys

import inquirer
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.chat_models import ChatOpenAI
from langchain.llms import AzureOpenAI, LlamaCpp, Ollama

from codeqai import utils
from codeqai.constants import LlmHost


class LLM:
    def __init__(self, llm_host: LlmHost, chat_model: str, deployment=None):
        if llm_host == LlmHost.OPENAI:
            self.chat_model = ChatOpenAI(
                temperature=0.9, max_tokens=2048, model=chat_model
            )
        elif llm_host == LlmHost.AZURE_OPENAI and deployment:
            self.chat_model = AzureOpenAI(
                temperature=0.9,
                max_tokens=2048,
                deployment_name=deployment,
                model=chat_model,
            )
        elif llm_host == LlmHost.LLAMACPP:
            self.install_llama_cpp()
            self.chat_model = LlamaCpp(
                model_path=chat_model,
                temperature=0.9,
                max_tokens=2048,
                verbose=False,
            )
        elif llm_host == LlmHost.OLLAMA:
            self.chat_model = Ollama(
                base_url="http://localhost:11434",
                model=chat_model,
                callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]),
            )

    def install_llama_cpp(self):
        try:
            from llama_cpp import Llama  # noqa: F401
        except ImportError:
            question = [
                inquirer.Confirm(
                    "confirm",
                    message=f"Local LLM interface package not found. Install {utils.get_bold_text('llama-cpp-python')}?",
                    default=True,
                ),
            ]

            answers = inquirer.prompt(question)
            if answers and answers["confirm"]:
                import platform

                def check_command(command):
                    try:
                        subprocess.run(
                            command,
                            check=True,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                        )
                        return True
                    except subprocess.CalledProcessError:
                        return False
                    except FileNotFoundError:
                        return False

                def install_llama(backend):
                    env_vars = {"FORCE_CMAKE": "1"}

                    if backend == "cuBLAS":
                        env_vars["CMAKE_ARGS"] = "-DLLAMA_CUBLAS=on"
                    elif backend == "hipBLAS":
                        env_vars["CMAKE_ARGS"] = "-DLLAMA_HIPBLAS=on"
                    elif backend == "Metal":
                        env_vars["CMAKE_ARGS"] = "-DLLAMA_METAL=on"
                    else:  # Default to OpenBLAS
                        env_vars[
                            "CMAKE_ARGS"
                        ] = "-DLLAMA_BLAS=ON -DLLAMA_BLAS_VENDOR=OpenBLAS"

                    try:
                        subprocess.run(
                            [
                                sys.executable,
                                "-m",
                                "pip",
                                "install",
                                "llama-cpp-python",
                            ],
                            env={**os.environ, **env_vars},
                            check=True,
                        )
                    except subprocess.CalledProcessError as e:
                        print(f"Error during installation with {backend}: {e}")

                def supports_metal():
                    # Check for macOS version
                    if platform.system() == "Darwin":
                        mac_version = tuple(map(int, platform.mac_ver()[0].split(".")))
                        # Metal requires macOS 10.11 or later
                        if mac_version >= (10, 11):
                            return True
                    return False

                # Check system capabilities
                if check_command(["nvidia-smi"]):
                    install_llama("cuBLAS")
                elif check_command(["rocminfo"]):
                    install_llama("hipBLAS")
                elif supports_metal():
                    install_llama("Metal")
                else:
                    install_llama("OpenBLAS")

                print("Finished downloading `Code-Llama` interface.")

                # Check if on macOS
                if platform.system() == "Darwin":
                    # Check if it's Apple Silicon
                    if platform.machine() != "arm64":
                        print(
                            "Warning: You are using Apple Silicon (M1/M2) Mac but your Python is not of 'arm64' architecture."
                        )
                        print(
                            "The llama.ccp x86 version will be 10x slower on Apple Silicon (M1/M2) Mac."
                        )
                        print(
                            "\nTo install the correct version of Python that supports 'arm64' architecture visit:"
                            "https://github.com/conda-forge/miniforge"
                        )

            else:
                exit("llama-cpp-python is required for local LLM.")

codeqai/llm.py
fynnfluegge/codeqai
import os

from git.repo import Repo


def get_git_root(path):
    git_repo = Repo(path, search_parent_directories=True)
    git_root = git_repo.git.rev_parse("--show-toplevel")
    return git_root


def load_files():
    git_root = get_git_root(os.getcwd())
    file_list = []

    for root, dirs, files in os.walk(git_root):
        if any(blacklist in root for blacklist in BLACKLIST_DIR):
            continue
        for file in files:
            file_ext = os.path.splitext(file)[1]
            if any(whitelist == file_ext for whitelist in WHITELIST_FILES):
                file_list.append(os.path.join(root, file))

    return file_list


BLACKLIST_DIR = [
    "__pycache__",
    ".venv",
    ".git",
    ".idea",
    "venv",
    "env",
    "node_modules",
    "dist",
    "build",
    ".vscode",
    ".github",
    ".gitlab",
    ".angular",
    "cdk.out",
    ".aws-sam",
    ".terraform",
]
WHITELIST_FILES = [
    ".js",
    ".mjs",
    ".ts",
    ".tsx",
    ".css",
    ".scss",
    ".less",
    ".html",
    ".htm",
    ".json",
    ".py",
    ".java",
    ".c",
    ".cpp",
    ".cs",
    ".go",
    ".php",
    ".rb",
    ".rs",
    ".swift",
    ".kt",
    ".scala",
    ".m",
    ".h",
    ".sh",
    ".pl",
    ".pm",
    ".lua",
    ".sql",
    ".yaml",
    ".yml",
    ".rst",
    ".md",
]

codeqai/repo.py
fynnfluegge/codeqai
from codeqai.treesitter.treesitter import (Treesitter,
                                                   TreesitterMethodNode)
from codeqai.treesitter.treesitter_c import TreesitterC
from codeqai.treesitter.treesitter_cpp import TreesitterCpp
from codeqai.treesitter.treesitter_cs import TreesitterCsharp
from codeqai.treesitter.treesitter_go import TreesitterGo
from codeqai.treesitter.treesitter_java import TreesitterJava
from codeqai.treesitter.treesitter_js import TreesitterJavascript
from codeqai.treesitter.treesitter_kt import TreesitterKotlin
from codeqai.treesitter.treesitter_py import TreesitterPython
from codeqai.treesitter.treesitter_rs import TreesitterRust
from codeqai.treesitter.treesitter_ts import TreesitterTypescript

codeqai/treesitter/__init__.py
fynnfluegge/codeqai
from abc import ABC

import tree_sitter
from tree_sitter_languages import get_language, get_parser

from codeqai.constants import Language
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterMethodNode:
    def __init__(
        self,
        name: "str | bytes | None",
        doc_comment: "str | None",
        node: tree_sitter.Node,
    ):
        self.name = name
        self.doc_comment = doc_comment
        self.method_source_code = node.text.decode()
        self.node = node


class Treesitter(ABC):
    def __init__(
        self,
        language: Language,
        method_declaration_identifier: str,
        name_identifier: str,
        doc_comment_identifier: str,
    ):
        self.parser = get_parser(language.value)
        self.language = get_language(language.value)
        self.method_declaration_identifier = method_declaration_identifier
        self.method_name_identifier = name_identifier
        self.doc_comment_identifier = doc_comment_identifier

    @staticmethod
    def create_treesitter(language: Language) -> "Treesitter":
        return TreesitterRegistry.create_treesitter(language)

    def parse(self, file_bytes: bytes) -> list[TreesitterMethodNode]:
        self.tree = self.parser.parse(file_bytes)
        result = []
        methods = self._query_all_methods(self.tree.root_node)
        for method in methods:
            method_name = self._query_method_name(method["method"])
            doc_comment = method["doc_comment"]
            result.append(
                TreesitterMethodNode(method_name, doc_comment, method["method"])
            )
        return result

    def _query_all_methods(
        self,
        node: tree_sitter.Node,
    ):
        methods = []
        if node.type == self.method_declaration_identifier:
            doc_comment_node = None
            if (
                node.prev_named_sibling
                and node.prev_named_sibling.type == self.doc_comment_identifier
            ):
                doc_comment_node = node.prev_named_sibling.text.decode()
            methods.append({"method": node, "doc_comment": doc_comment_node})
        else:
            for child in node.children:
                methods.extend(self._query_all_methods(child))
        return methods

    def _query_method_name(self, node: tree_sitter.Node):
        if node.type == self.method_declaration_identifier:
            for child in node.children:
                if child.type == self.method_name_identifier:
                    return child.text.decode()
        return None

codeqai/treesitter/treesitter.py
fynnfluegge/codeqai
import tree_sitter

from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterC(Treesitter):
    def __init__(self):
        super().__init__(Language.C, "function_definition", "identifier", "comment")

    def _query_method_name(self, node: tree_sitter.Node):
        if node.type == self.method_declaration_identifier:
            for child in node.children:
                # if method returns pointer, skip pointer declarator
                if child.type == "pointer_declarator":
                    child = child.children[1]
                if child.type == "function_declarator":
                    for child in child.children:
                        if child.type == self.method_name_identifier:
                            return child.text.decode()
        return None


# Register the TreesitterJava class in the registry
TreesitterRegistry.register_treesitter(Language.C, TreesitterC)

codeqai/treesitter/treesitter_c.py
fynnfluegge/codeqai
import tree_sitter

from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterCpp(Treesitter):
    def __init__(self):
        super().__init__(Language.CPP, "function_definition", "identifier", "comment")

    def _query_method_name(self, node: tree_sitter.Node):
        if node.type == self.method_declaration_identifier:
            for child in node.children:
                # if method returns pointer, skip pointer declarator
                if child.type == "pointer_declarator":
                    child = child.children[1]
                if child.type == "function_declarator":
                    for child in child.children:
                        if child.type == self.method_name_identifier:
                            return child.text.decode()
        return None


# Register the TreesitterJava class in the registry
TreesitterRegistry.register_treesitter(Language.CPP, TreesitterCpp)

codeqai/treesitter/treesitter_cpp.py
fynnfluegge/codeqai
import tree_sitter

from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterCsharp(Treesitter):
    def __init__(self):
        super().__init__(
            Language.C_SHARP, "method_declaration", "identifier", "comment"
        )

    def _query_method_name(self, node: tree_sitter.Node):
        first_match = None
        if node.type == self.method_declaration_identifier:
            for child in node.children:
                # if the return type is an object type, then the method name
                # is the second match
                if child.type == self.method_name_identifier and not first_match:
                    first_match = child.text.decode()
                elif child.type == self.method_name_identifier and first_match:
                    return child.text.decode()
        return first_match

    def _query_all_methods(self, node: tree_sitter.Node):
        methods = []
        if node.type == self.method_declaration_identifier:
            doc_comment_nodes = []
            if (
                node.prev_named_sibling
                and node.prev_named_sibling.type == self.doc_comment_identifier
            ):
                current_doc_comment_node = node.prev_named_sibling
                while (
                    current_doc_comment_node
                    and current_doc_comment_node.type == self.doc_comment_identifier
                ):
                    doc_comment_nodes.append(current_doc_comment_node.text.decode())
                    if current_doc_comment_node.prev_named_sibling:
                        current_doc_comment_node = (
                            current_doc_comment_node.prev_named_sibling
                        )
                    else:
                        current_doc_comment_node = None

            doc_comment_str = ""
            doc_comment_nodes.reverse()
            for doc_comment_node in doc_comment_nodes:
                doc_comment_str += doc_comment_node + "\n"
            if doc_comment_str.strip() != "":
                methods.append({"method": node, "doc_comment": doc_comment_str.strip()})
            else:
                methods.append({"method": node, "doc_comment": None})
        else:
            for child in node.children:
                methods.extend(self._query_all_methods(child))
        return methods


# Register the TreesitterJava class in the registry
TreesitterRegistry.register_treesitter(Language.C_SHARP, TreesitterCsharp)

codeqai/treesitter/treesitter_cs.py
fynnfluegge/codeqai
from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterGo(Treesitter):
    def __init__(self):
        super().__init__(Language.GO, "function_declaration", "identifier", "comment")


# Register the TreesitterJava class in the registry
TreesitterRegistry.register_treesitter(Language.GO, TreesitterGo)

codeqai/treesitter/treesitter_go.py
fynnfluegge/codeqai
from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterJava(Treesitter):
    def __init__(self):
        super().__init__(
            Language.JAVA, "method_declaration", "identifier", "block_comment"
        )


# Register the TreesitterJava class in the registry
TreesitterRegistry.register_treesitter(Language.JAVA, TreesitterJava)

codeqai/treesitter/treesitter_java.py
fynnfluegge/codeqai
from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterJavascript(Treesitter):
    def __init__(self):
        super().__init__(
            Language.JAVASCRIPT, "function_declaration", "identifier", "comment"
        )


# Register the TreesitterJava class in the registry
TreesitterRegistry.register_treesitter(Language.JAVASCRIPT, TreesitterJavascript)

codeqai/treesitter/treesitter_js.py
fynnfluegge/codeqai
from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterKotlin(Treesitter):
    def __init__(self):
        super().__init__(
            Language.KOTLIN, "function_declaration", "simple_identifier", "comment"
        )


# Register the TreesitterJava class in the registry
TreesitterRegistry.register_treesitter(Language.KOTLIN, TreesitterKotlin)

codeqai/treesitter/treesitter_kt.py
fynnfluegge/codeqai
import tree_sitter

from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter, TreesitterMethodNode
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterPython(Treesitter):
    def __init__(self):
        super().__init__(
            Language.PYTHON, "function_definition", "identifier", "expression_statement"
        )

    def parse(self, file_bytes: bytes) -> list[TreesitterMethodNode]:
        self.tree = self.parser.parse(file_bytes)
        result = []
        methods = self._query_all_methods(self.tree.root_node)
        for method in methods:
            method_name = self._query_method_name(method)
            doc_comment = self._query_doc_comment(method)
            result.append(TreesitterMethodNode(method_name, doc_comment, method))
        return result

    def _query_method_name(self, node: tree_sitter.Node):
        if node.type == self.method_declaration_identifier:
            for child in node.children:
                if child.type == self.method_name_identifier:
                    return child.text.decode()
        return None

    def _query_all_methods(self, node: tree_sitter.Node):
        methods = []
        for child in node.children:
            if child.type == self.method_declaration_identifier:
                methods.append(child)
        return methods

    def _query_doc_comment(self, node: tree_sitter.Node):
        query_code = """
            (function_definition
                body: (block . (expression_statement (string)) @function_doc_str))
        """
        doc_str_query = self.language.query(query_code)
        doc_strs = doc_str_query.captures(node)

        if doc_strs:
            return doc_strs[0][0].text.decode()
        else:
            return None


# Register the TreesitterPython class in the registry
TreesitterRegistry.register_treesitter(Language.PYTHON, TreesitterPython)

codeqai/treesitter/treesitter_py.py
fynnfluegge/codeqai
from codeqai.constants import Language


class TreesitterRegistry:
    _registry = {}

    @classmethod
    def register_treesitter(cls, name, treesitter_class):
        cls._registry[name] = treesitter_class

    @classmethod
    def create_treesitter(cls, name: Language):
        treesitter_class = cls._registry.get(name)
        if treesitter_class:
            return treesitter_class()
        else:
            raise ValueError("Invalid tree type")

codeqai/treesitter/treesitter_registry.py
fynnfluegge/codeqai
import tree_sitter

from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterRust(Treesitter):
    def __init__(self):
        super().__init__(Language.RUST, "function_item", "identifier", "line_comment")

    def _query_all_methods(self, node: tree_sitter.Node):
        methods = []
        if node.type == self.method_declaration_identifier:
            doc_comment_nodes = []
            if (
                node.prev_named_sibling
                and node.prev_named_sibling.type == self.doc_comment_identifier
            ):
                current_doc_comment_node = node.prev_named_sibling
                while (
                    current_doc_comment_node
                    and current_doc_comment_node.type == self.doc_comment_identifier
                ):
                    doc_comment_nodes.append(current_doc_comment_node.text.decode())
                    if current_doc_comment_node.prev_named_sibling:
                        current_doc_comment_node = (
                            current_doc_comment_node.prev_named_sibling
                        )
                    else:
                        current_doc_comment_node = None

            doc_comment_str = ""
            doc_comment_nodes.reverse()
            for doc_comment_node in doc_comment_nodes:
                doc_comment_str += doc_comment_node + "\n"
            if doc_comment_str.strip() != "":
                methods.append({"method": node, "doc_comment": doc_comment_str.strip()})
            else:
                methods.append({"method": node, "doc_comment": None})
        else:
            for child in node.children:
                methods.extend(self._query_all_methods(child))
        return methods


# Register the TreesitterJava class in the registry
TreesitterRegistry.register_treesitter(Language.RUST, TreesitterRust)

codeqai/treesitter/treesitter_rs.py
fynnfluegge/codeqai
from codeqai.constants import Language
from codeqai.treesitter.treesitter import Treesitter
from codeqai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterTypescript(Treesitter):
    def __init__(self):
        super().__init__(
            Language.TYPESCRIPT, "function_declaration", "identifier", "comment"
        )


# Register the TreesitterJava class in the registry
TreesitterRegistry.register_treesitter(Language.TYPESCRIPT, TreesitterTypescript)

codeqai/treesitter/treesitter_ts.py
fynnfluegge/codeqai
import os

import langchain.text_splitter as text_splitter

from codeqai.constants import Language


def get_programming_language(file_extension: str) -> Language:
    """
    Returns the programming language based on the provided file extension.

    Args:
        file_extension (str): The file extension to determine the programming language of.

    Returns:
        Language: The programming language corresponding to the file extension. If the file extension is not found
        in the language mapping, returns Language.UNKNOWN.
    """
    language_mapping = {
        ".py": Language.PYTHON,
        ".js": Language.JAVASCRIPT,
        ".jsx": Language.JAVASCRIPT,
        ".mjs": Language.JAVASCRIPT,
        ".cjs": Language.JAVASCRIPT,
        ".ts": Language.TYPESCRIPT,
        ".tsx": Language.TYPESCRIPT,
        ".java": Language.JAVA,
        ".kt": Language.KOTLIN,
        ".rs": Language.RUST,
        ".go": Language.GO,
        ".cpp": Language.CPP,
        ".c": Language.C,
        ".cs": Language.C_SHARP,
    }
    return language_mapping.get(file_extension, Language.UNKNOWN)


def get_file_extension(file_name: str) -> str:
    """
    Returns the extension of a file from its given name.

    Parameters:
        file_name (str): The name of the file.

    Returns:
        str: The extension of the file.

    """
    return os.path.splitext(file_name)[-1]


def get_langchain_language(language: Language) -> text_splitter.Language | None:
    if language == Language.PYTHON:
        return text_splitter.Language.PYTHON
    elif language == Language.JAVASCRIPT:
        return text_splitter.Language.JS
    elif language == Language.TYPESCRIPT:
        return text_splitter.Language.TS
    elif language == Language.JAVA:
        return text_splitter.Language.JAVA
    elif language == Language.KOTLIN:
        return text_splitter.Language.KOTLIN
    elif language == Language.RUST:
        return text_splitter.Language.RUST
    elif language == Language.GO:
        return text_splitter.Language.GO
    elif language == Language.CPP:
        return text_splitter.Language.CPP
    elif language == Language.C_SHARP:
        return text_splitter.Language.CSHARP

    return None


def get_bold_text(text):
    return f"\033[01m{text}\033[0m"

codeqai/utils.py
fynnfluegge/codeqai
import inquirer
from langchain.embeddings.base import Embeddings
from langchain.schema import Document
from langchain.vectorstores import FAISS
from yaspin import yaspin

from codeqai import utils
from codeqai.config import get_cache_path


class VectorStore:
    def __init__(
        self, name: str, embeddings: Embeddings, documents: list[Document] = []
    ):
        self.name = name
        self.embeddings = embeddings
        self.install_faiss()
        if not documents:
            spinner = yaspin(text="💾 Loading vector store...", color="green")
            spinner.start()
            self.db = FAISS.load_local(
                index_name=self.name,
                folder_path=get_cache_path(),
                embeddings=self.embeddings,
            )
            spinner.stop()
        else:
            spinner = yaspin(text="💾 Indexing vector store...", color="green")
            spinner.start()
            self.db = FAISS.from_documents(documents, embeddings)
            self.db.save_local(index_name=self.name, folder_path=get_cache_path())
            spinner.stop()
        self.retriever = self.db.as_retriever(search_type="mmr", search_kwargs={"k": 8})

    def similarity_search(self, query: str):
        return self.db.similarity_search(query, k=4)

    def install_faiss(self):
        try:
            import faiss  # noqa: F401
        except ImportError:
            question = [
                inquirer.Confirm(
                    "confirm",
                    message=f"{utils.get_bold_text('faiss')} package not found in this python environment. Do you want to install it now?",
                    default=True,
                ),
            ]

            answers = inquirer.prompt(question)
            if answers and answers["confirm"]:
                import subprocess
                import sys

                question = [
                    inquirer.List(
                        "faiss-installation",
                        message=f"Please select the appropriate option to install FAISS.",
                        choices=[
                            "faiss-cpu",
                            "faiss-gpu (Only if your system supports CUDA))",
                        ],
                        default="faiss-cpu",
                    ),
                ]

                answers = inquirer.prompt(question)
                if answers and answers["faiss-installation"]:
                    try:
                        subprocess.run(
                            [
                                sys.executable,
                                "-m",
                                "pip",
                                "install",
                                answers["faiss-installation"],
                            ],
                            check=True,
                        )
                    except subprocess.CalledProcessError as e:
                        print(f"Error during faiss installation: {e}")
            else:
                exit("faiss package is required for codeqai to work.")

codeqai/vector_store.py
fynnfluegge/codeqai