ML-engineer-challenge.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Background:\n",
    "You are provided with a boilerplate Python notebook that interacts with a pull request (PR) reviewer system using a Large Language Model (LLM). This system is designed to review code by identifying potential issues in the changes made in a PR. The initial setup involves parsing code from the repository, storing code snippets in a vector database (using ChromaDB), and querying this database to provide contextually relevant code examples for LLM-based code review. The reviews are done file-by-file with prompts structured to guide LLM in generating constructive feedback in a specific JSON format.\n",
    "\n",
    "## Task Overview:\n",
    "Your challenge is to conceptualize and implement enhancements in three key areas of the existing PR reviewer system:\n",
    "1. **Splitting files into chunks**\n",
    "2. **Retriever query generation system**\n",
    "3. **LLM review prompt structure**\n",
    "\n",
    "\n",
    "The challenge focuses on three critical areas requiring innovation and enhancement:\n",
    "\n",
    "1. **Splitting Files into Chunks**: You have to develop a sophisticated method for dividing repository code into meaningfully sized chunks or propose an entirely new approach for managing code snippets. The goal is to enhance logical or functional coherence within the chunks without losing efficiency.\n",
    "\n",
    "2. **Retriever Query Generator System**: You have to refine the system’s approach to generating queries for retrieving relevant documents from a vector database. You should aim to improve the precision and relevance of document retrieval and might want to explore advanced query generation strategies and incorporating richer contextual information.\n",
    "\n",
    "3. **LLM Review Prompt Revamping**: Currently, the system reviews code on a per-file basis. You are encouraged to explore alternative methodologies, such as reviewing entire PRs in a holistic manner or employing a sophisticated LLM agent system for more effective and context-aware reviews.\n",
    "\n",
    "**You may choose to work on just one or any number of the above points.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv(\".env\")\n",
    "\n",
    "assert os.getenv(\"GITHUB_TOKEN\") is not None and os.getenv(\"GITHUB_TOKEN\") != \"a_key_goes_here\"\n",
    "assert os.getenv(\"OPENAI_API_KEY\") is not None and os.getenv(\"OPENAI_API_KEY\") != \"a_key_goes_here\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to Pre-made Tools\n",
    "\n",
    "In order to streamline the process of interacting with and analyzing repository files for our PR reviewer system, we've developed a suite of pre-made tools. These utilities are designed to handle various aspects of file manipulation, diff extraction, and repository management seamlessly. Below is an overview of each tool and its primary purpose.\n",
    "\n",
    "### LocalFile\n",
    "\n",
    "The `LocalFile` class acts as a refined interface for file interaction, offering a simplified way to access a file's contents along with pertinent metadata. It's an enhanced representation that abstracts away the complexities of raw file handling, making file operations more intuitive and less error-prone.\n",
    "\n",
    "### DiffRepresentation\n",
    "\n",
    "Dealing with diffs can be intricate, given their crucial role in identifying changes between file versions. The `DiffRepresentation` class simplifies this task by providing methods to extract and analyze diff metadata effectively. It specifically aids in parsing edit locations within a file, such as the line numbers of changed sections and the position of the [diff hunk header](https://stackoverflow.com/q/28111035), which marks the start of a set of differences in file content. This specialized tool ensures that diffs are handled accurately, facilitating a deeper analysis of code changes.\n",
    "\n",
    "### GithubUtils\n",
    "\n",
    "The `GithubUtils` class encapsulates functionality for fetching information and files from a specific GitHub repository and its associated PRs. It leverages the PyGithub library to abstract away the intricacies of API communication, offering a more user-friendly interface for retrieving the data needed for review analysis.\n",
    "\n",
    "### LocalRepository\n",
    "\n",
    "To further mitigate the limitations posed by GitHub's rate limiting and to enhance the efficiency of repository analysis, the `LocalRepository` class provides a mechanism to mirror a GitHub repository locally.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import github\n",
    "from korbit_tools.github_service import GithubUtils\n",
    "\n",
    "rest_github = github.Github(os.getenv(\"GITHUB_TOKEN\"))\n",
    "\n",
    "REPOSITORY = \"langchain-ai/langchain\"\n",
    "PR_NUMBER = 13999\n",
    "\n",
    "ALLOWED_EXTENSIONS = [\".py\"]\n",
    "\n",
    "repo = rest_github.get_repo(REPOSITORY)\n",
    "pr = repo.get_pull(PR_NUMBER)\n",
    "\n",
    "# Here is how we get content for a Pull Request\n",
    "for content_file, pr_diff in GithubUtils.get_pull_request_content_file_iter(repo, pr, allowed_extensions=ALLOWED_EXTENSIONS):\n",
    "    print(pr_diff.diff[:100])\n",
    "    print(content_file.contents[:100])\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vector store\n",
    "### Download repository locally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from korbit_tools.github_service import download_repository\n",
    "from korbit_tools.repository_search import LocalRepository\n",
    "\n",
    "\n",
    "repo_path = download_repository(repo, pr.head.sha, \"./repositories\")\n",
    "local_repository = LocalRepository(repo_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load local repository\n",
    "\n",
    "Here we convert the repository's files, into langchain Documents model, to then store them into ChromDB.\n",
    "\n",
    "[https://python.langchain.com/docs/integrations/document_loaders/source_code#splitting](https://python.langchain.com/docs/integrations/document_loaders/source_code#splitting)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import PythonLoader\n",
    "from langchain_community.document_loaders import DirectoryLoader\n",
    "\n",
    "loader = DirectoryLoader(repo_path,glob=\"**/*.py\", loader_cls=PythonLoader)\n",
    "\n",
    "documents = loader.load()\n",
    "len(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Split files into chunks\n",
    "\n",
    "We need to chunk the files into multiple ones otherwise the context will have too much context and the interesting information will be lost in the context.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter, Language\n",
    "\n",
    "py_splitter = RecursiveCharacterTextSplitter.from_language(\n",
    "    language=Language.PYTHON, chunk_size=3000, chunk_overlap=100\n",
    ")\n",
    "texts = py_splitter.split_documents(documents)\n",
    "len(texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell will take several minutes to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "# The vector store will be persisted in the current directory as .chromadb folder\n",
    "db = Chroma.from_documents(texts, OpenAIEmbeddings(disallowed_special=()), persist_directory=\".chromadb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "# This cell allow you to retrieve the vector store from the previously created one\n",
    "db = Chroma(persist_directory=\".chromadb\", embedding_function=OpenAIEmbeddings(disallowed_special=()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retriever\n",
    "Here is the configuration of the vectorDB retriever. You can obviously change that to match what you think will work best. For example using a subset of the diff, or the all content of the file you are reviewing.\n",
    "You can use this function repository_search to query the all vectorial database, containing embeddings of code snippets from the repository the PR request is made on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import textwrap\n",
    "\n",
    "\n",
    "retriever = db.as_retriever(\n",
    "    search_type=\"mmr\", # https://api.python.langchain.com/en/latest/retrievers/langchain_community.retrievers.zep.SearchType.html\n",
    "    search_kwargs={\"k\": 4},\n",
    ")\n",
    "\n",
    "INDENTATION = \"  \"\n",
    "def repository_search(query: str) -> str:\n",
    "    relevant_docs = retriever.get_relevant_documents(query)\n",
    "    output = \"\"\n",
    "    for doc in relevant_docs:\n",
    "        output += f\"- {doc.metadata[\"source\"]}\\n{INDENTATION}```{doc.metadata.get(\"language\", \"py\")}\\n{textwrap.indent(doc.page_content, INDENTATION)}\\n{INDENTATION}```\\n\\n\"\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prompt\n",
    "We are providing a boilerplate prompt for reviewing one file at a time, but feel free to customize it to your liking. Modify the steps, inputs, outputs, and overall prompt as you see fit to make it your own.\n",
    "\n",
    "By default, the expected output for this prompt is in JSON format as shown below. However, you are welcome to alter this structure if you believe there is a more effective way to provide feedback:\n",
    "\n",
    "```json\n",
    "[\n",
    "    {\n",
    "        \"description\": \"A brief description of the issue.\",\n",
    "        \"code_snippet\": \"The relevant code snippet causing the issue.\",\n",
    "        \"category\": \"The category of the issue.\",\n",
    "        \"severity\": \"Severity level of the detected issue on a scale from 1 to 10.\"\n",
    "    }\n",
    "]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "\n",
    "ISSUE_DETECTOR_SYSTEM = \"\"\"\\\n",
    "You are a senior software engineer tasked with reviewing a pull request, a set of proposed changes to the code base, submitted by one of the software engineers on your team. Your goal is to identify issues in the pull request and provide constructive feedback to help improve the quality of the code base.\n",
    "\n",
    "The review must be in JSON format and must contain a list of the following attributes for each issue detected:\n",
    "1. description: Constructive feedback that explains the issue in the pull request.\n",
    "2. code_snippet: The code snippet the description refers to.\n",
    "3. category: The type of issue detected: bug, good practices, or other.\n",
    "4. severity: The urgency of the detected issue. Scale from 1 to 10.\n",
    "\n",
    "Based on the above context, you must output a JSON list of issues that you have identified. For each issue, use the following attributes:\n",
    "```json\n",
    "[{{\n",
    "  \"description\": \"An issue description you found in the diff\",\n",
    "  \"code_snippet\": \"def my_function(x):\\n    return x**2\\n\\nmy_function(123)\",\n",
    "  \"category\": \"naming\",\n",
    "  \"severity\": 1\n",
    "}}]\n",
    "```\n",
    "\n",
    "If you didn't find any issues, output an empty list, but you must output something.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "ISSUE_DETECTOR_AGENT_HUMAN = \"\"\"\n",
    "- title: {pr_title}\n",
    "- description: {pr_description}\n",
    "- file path: {file_path}\n",
    "- file content:\n",
    "\n",
    "```\n",
    "{content_file}\n",
    "```\n",
    "\n",
    "- diff representation:\n",
    "\n",
    "```diff\n",
    "{pr_diff}\n",
    "```\n",
    "\n",
    "- additional context useful for the pull request review:\n",
    "```\n",
    "{context}\n",
    "```\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "prompt_reviewer = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", ISSUE_DETECTOR_SYSTEM),\n",
    "        (\"human\", ISSUE_DETECTOR_AGENT_HUMAN)\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from korbit_tools.local_file import LocalFile\n",
    "from korbit_tools.diff_representation import DiffRepresentation\n",
    "\n",
    "\n",
    "def compute_context(content_file: LocalFile, pr_diff: DiffRepresentation) -> str:\n",
    "    \"\"\"\n",
    "    This is a very basic version of the RAG system\n",
    "    that returns the most relevant context for the pull request from the vector store.\n",
    "    \"\"\"\n",
    "    query = pr_diff.diff\n",
    "    output = repository_search(query)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PR review\n",
    "\n",
    "After setting up our system to analyze pull requests using the Large Language Model (LLM), we now enter the critical phase of putting our setup to work. In essence, our goal is to systematically review each file associated with a given pull request. This is achieved by leveraging our vector database to unearth relevant code snippets and differences, which then become the foundation for our LLM-based code review.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import LLMChain\n",
    "from langchain_openai.chat_models import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4-1106-preview\")\n",
    "chain = LLMChain(llm=llm, prompt=prompt_reviewer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "from korbit_tools.llm_utils import count_token_string\n",
    "from korbit_tools.string_search import extract_json_from_text\n",
    "\n",
    "# NOTE: We won't review files that are above 60k tokens long. The model won't be able to process it along other files retrieved by the similarity search.\n",
    "CONTENT_FILE_TOKEN_LIMIT = 60000\n",
    "\n",
    "pr_issues = {}\n",
    "llm_inputs = {\"pr_title\": pr.title, \"pr_description\": pr.body}\n",
    "# Get all files that have been changed in the Pull Request\n",
    "for content_file, pr_diff in GithubUtils.get_pull_request_content_file_iter(\n",
    "    repo, pr, allowed_extensions=ALLOWED_EXTENSIONS\n",
    "):\n",
    "\n",
    "    token_count = count_token_string(content_file.contents)\n",
    "    if token_count >= CONTENT_FILE_TOKEN_LIMIT:\n",
    "        continue\n",
    "\n",
    "    # Setup the LLM inputs to match the variables in the prompt above\n",
    "    llm_inputs[\"content_file\"] = content_file.contents\n",
    "    llm_inputs[\"file_path\"] = content_file.path\n",
    "    llm_inputs[\"pr_diff\"] = pr_diff.diff\n",
    "    llm_inputs[\"context\"] = compute_context(content_file, pr_diff)\n",
    "\n",
    "    output = chain.invoke(llm_inputs)\n",
    "    \n",
    "    pr_issues[content_file.path] = output.get(\"text\")\n",
    "\n",
    "    # Remove this break statement to review all files in the PR\n",
    "    break\n",
    "\n",
    "for file_path, raw_issues in pr_issues.items():\n",
    "    pprint(extract_json_from_text(raw_issues))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cleanup downloaded repositories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "!rm -rf ./repositories/* --exclude='.gitkeep'"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "code-analysis-worker",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}