notebooks/Generate-and-Ingest-Documentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3af5dbf3",
   "metadata": {},
   "source": [
    "# Code Analyis using Amazon Q for Business\n",
    "\n",
    "In this notebook we will use Amazon Q for Business to analyze the code of a repository. We will use langchain-agents repository in `aws-samples` as an example."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6c11e19",
   "metadata": {},
   "source": [
    "## Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e24b3d4",
   "metadata": {},
   "source": [
    "First make sure that your AWS credentials are properly configured. You can do this by installing the AWS CLI and running `aws configure`:\n",
    "\n",
    "```bash\n",
    "pip install awscli\n",
    "aws configure\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cf5b2e1",
   "metadata": {},
   "source": [
    "Then install the necessary libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c2ca07a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install boto3 --upgrade \n",
    "%pip install GitPython shutils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bf81184",
   "metadata": {},
   "source": [
    "## Creating the Amazon Q Application"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "172f4c48",
   "metadata": {},
   "source": [
    "In the steps below we will create an Amazon Q application that will be used to process and then answer questions about the code of a repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2a09214",
   "metadata": {},
   "outputs": [],
   "source": [
    "amazon_q_user_id = \"<your email>\"\n",
    "role_arn = None\n",
    "amazon_q_app_id = None\n",
    "reuse_existing_q_app = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a5242ce",
   "metadata": {},
   "source": [
    "If you want to reuse the application, uncomment the cell below, fill in the values, and run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8468d8c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# amazon_q_app_id = <your app id>\n",
    "# role_arn = <your role arn>\n",
    "# reuse_existing_q_app = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad04fa07",
   "metadata": {},
   "source": [
    "Next we create the IAM role that will be used by the Amazon Q application to access the repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a31bb5e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Q IAM Service role from iam-policy.json and trust-policy.json\n",
    "import boto3\n",
    "import json\n",
    "import datetime\n",
    "\n",
    "if not role_arn:\n",
    "\n",
    "    # Create IAM role\n",
    "    project_name = f\"Code-Analysis-Demo-App-{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}\"\n",
    "    iam = boto3.client('iam')\n",
    "    # Note to work with Q Business the role MUST start with \"QBusiness-Application-\"\n",
    "    role_name = f\"QBusiness-Application-{project_name}\"\n",
    "    role_policy_file = \"./security/iam-policy.json\"\n",
    "    trust_policy_file = \"./security/trust-policy.json\"\n",
    "\n",
    "    # Create role using iam policy and trust policy\n",
    "    role_policy = json.load(open(role_policy_file))\n",
    "    trust_policy = json.load(open(trust_policy_file))\n",
    "    role = iam.create_role(\n",
    "        RoleName=role_name,\n",
    "        AssumeRolePolicyDocument=json.dumps(trust_policy)\n",
    "    )\n",
    "    role_arn = role['Role']['Arn']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e306bb1a",
   "metadata": {},
   "source": [
    "Now we create the Amazon Q application passing the role created above. Note that we are enabling attachments so that we can later process every file of code in the repository with the application"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bcfc678",
   "metadata": {},
   "outputs": [],
   "source": [
    "amazon_q = boto3.client('qbusiness')\n",
    "\n",
    "if not reuse_existing_q_app:\n",
    "    response = amazon_q.create_application(\n",
    "        attachmentsConfiguration={\n",
    "            'attachmentsControlMode': 'ENABLED'\n",
    "        },\n",
    "        description=f\"{project_name}-{datetime.datetime.now().strftime('%Y-%m-%d')}\",\n",
    "        displayName=project_name,\n",
    "        roleArn=role_arn,\n",
    "    )\n",
    "    amazon_q_app_id = response[\"applicationId\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b758dfdb",
   "metadata": {},
   "source": [
    "After creating the application, we will create an index that will be used to store the information about the code of the repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc06cf89",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not reuse_existing_q_app:\n",
    "    response = amazon_q.create_index(\n",
    "        applicationId=amazon_q_app_id,\n",
    "        capacityConfiguration={\n",
    "            'units': 1\n",
    "        },\n",
    "        description=f\"{project_name}-{datetime.datetime.now().strftime('%Y-%m-%d')}\",\n",
    "        displayName=project_name,\n",
    "    )\n",
    "    index_id = response[\"indexId\"]\n",
    "else:\n",
    "    response = amazon_q.list_indices(\n",
    "        applicationId=amazon_q_app_id\n",
    "    )\n",
    "    index_id = response[\"indices\"][0]['indexId']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10670420",
   "metadata": {},
   "source": [
    "Lastly we will create a retriever to fetch the relevant information when we ask questions about the repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27ce07f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not reuse_existing_q_app:\n",
    "    response = amazon_q.create_retriever(\n",
    "        applicationId=amazon_q_app_id,\n",
    "        configuration={\n",
    "            'nativeIndexConfiguration': {\n",
    "                'indexId': index_id\n",
    "            }\n",
    "        },\n",
    "        displayName=project_name,\n",
    "        roleArn=role_arn,\n",
    "        type='NATIVE_INDEX'\n",
    "    )\n",
    "    retriever_id = response[\"retrieverId\"]\n",
    "else:\n",
    "    retriever_id = amazon_q.list_retrievers(\n",
    "        maxResults=1,\n",
    "        applicationId=amazon_q_app_id,\n",
    "    )[\"retrievers\"][0][\"retrieverId\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b78f1338",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "while True:\n",
    "    response = amazon_q.get_index(\n",
    "    applicationId=amazon_q_app_id,\n",
    "    indexId=index_id,\n",
    "    )\n",
    "    status = response.get('status')\n",
    "    print(f\"Creat index status {status}\")\n",
    "    if status == 'ACTIVE':\n",
    "        break\n",
    "    time.sleep(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76b49e0b",
   "metadata": {},
   "source": [
    "## Generating and Ingesting Documentation\n",
    "If we only ingest the code, we will be retrieving random code chunks that may not be relevant to the questions we want to ask. To avoid this, we will generate concise documentation for the repository and ingest it into the index."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2713a327",
   "metadata": {},
   "source": [
    "First we will define a few helper functions. The first one will be used to take a file and a prompt and return the answer generated from Q."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "355db82d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ask_question_with_attachment(prompt, filename):\n",
    "    data=open(filename, 'rb')\n",
    "    answer = amazon_q.chat_sync(\n",
    "        applicationId=amazon_q_app_id,\n",
    "        userId=amazon_q_user_id,\n",
    "        userMessage=prompt,\n",
    "        attachments=[\n",
    "            {\n",
    "                'data': data.read(),\n",
    "                'name': filename\n",
    "            },\n",
    "        ],\n",
    "    )\n",
    "    return answer['systemMessage']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cacad47e",
   "metadata": {},
   "source": [
    "Next, we define a function that will upload the generated answer, the filename, and the prompt to the index along with a source attribute that can be used later to find the files that the answer came from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d07cac04",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "def upload_prompt_answer_and_file_name(filename, prompt, answer, repo_url):\n",
    "    amazon_q.batch_put_document(\n",
    "        applicationId=amazon_q_app_id,\n",
    "        indexId=index_id,\n",
    "        roleArn=role_arn,\n",
    "        documents=[\n",
    "            {\n",
    "                'id': str(uuid.uuid4()),\n",
    "                'contentType': 'PLAIN_TEXT',\n",
    "                'title': filename,\n",
    "                'content':{\n",
    "                    'blob': f\"{filename} | {prompt} | {answer}\".encode('utf-8')\n",
    "                },\n",
    "                'attributes': [\n",
    "                    {\n",
    "                        'name': 'url',\n",
    "                        'value': {\n",
    "                            'stringValue': f\"{repo_url}{filename}\"\n",
    "                        }\n",
    "                    }\n",
    "                ]\n",
    "            },\n",
    "        ]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8a3734e",
   "metadata": {},
   "source": [
    "We also create a helper function to save the generated answers to a local folder in case we want to inspect it later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b42ffcae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to save generated answers to folder documentation/\n",
    "def save_answers(answer, filepath, folder):\n",
    "    import os\n",
    "    # Only create directory until the last / of filepath\n",
    "    sub_directory = f\"{folder}{filepath[:filepath.rfind('/')+1]}\"\n",
    "    if not os.path.exists(sub_directory):\n",
    "        # Only create directory until the last /\n",
    "        os.makedirs(sub_directory)\n",
    "    # Replace all file endings with .txt\n",
    "    filepath = filepath[:filepath.rfind('.')] + \".txt\"\n",
    "    with open(f\"{folder}{filepath}\", \"w\") as f:\n",
    "        f.write(answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f16aaef2",
   "metadata": {},
   "source": [
    "Since there are some files we want to ignore, we will define a function to filter the files we want to process. You can modify this function to fit your needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9ad89002",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "def should_ignore_path(path):\n",
    "    path_components = path.split(os.sep)\n",
    "    for component in path_components:\n",
    "        if component.startswith('.'):\n",
    "            return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3716b194",
   "metadata": {},
   "source": [
    "## Processing and Ingesting a Git Repository\n",
    "\n",
    "Now that we have the application and the index, along with the helper functions, we can process every file in the repository and ingest the generated answers into the index.\n",
    "\n",
    "First we will clone the repository and then process every file in the repository. We will use the helper functions to generate the answers and ingest them into the index. We will also save the generated answers to a local folder.\n",
    "\n",
    "If you want to create documentation with focus on a particular aspect, i.e. security, you can change the prompt so Q will generate answers with that focus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a33cf263",
   "metadata": {},
   "outputs": [],
   "source": [
    "import git\n",
    "import shutil\n",
    "\n",
    "def process_repository(repo_url, ssh_url=None):\n",
    "\n",
    "    # Temporary clone location\n",
    "    tmp_dir = f\"/tmp/{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}\" \n",
    "\n",
    "    destination_folder = 'repositories/'\n",
    "\n",
    "    # Clone the repository\n",
    "    # If you authenticate with some other repo provider just change the line below\n",
    "    if ssh_url:\n",
    "        repo = git.Repo.clone_from(ssh_url, tmp_dir)\n",
    "    else:\n",
    "        repo = git.Repo.clone_from(repo_url, tmp_dir)\n",
    "\n",
    "    # Copy all files to destination folder\n",
    "    for src_dir, dirs, files in os.walk(tmp_dir):\n",
    "        dst_dir = src_dir.replace(tmp_dir, destination_folder)\n",
    "        if not os.path.exists(dst_dir):\n",
    "            os.mkdir(dst_dir)\n",
    "        for file_ in files:\n",
    "            src_file = os.path.join(src_dir, file_)\n",
    "            dst_file = os.path.join(dst_dir, file_)\n",
    "            if os.path.exists(dst_file):\n",
    "                os.remove(dst_file)\n",
    "            shutil.copy(src_file, dst_dir)\n",
    "    \n",
    "    # Delete temp clone       \n",
    "    shutil.rmtree(tmp_dir)\n",
    "\n",
    "    import time\n",
    "\n",
    "    processed_files = []\n",
    "    failed_files = []\n",
    "\n",
    "    for root, dirs, files in os.walk(destination_folder):\n",
    "        if should_ignore_path(root):\n",
    "            continue\n",
    "        for file in files:\n",
    "            if file.endswith(('.png', '.jpg', '.jpeg', '.gif', '.zip')):\n",
    "                continue\n",
    "            # Ignore files that start with a dot (.)\n",
    "            if file.startswith('.'):\n",
    "                continue\n",
    "                \n",
    "            file_path = os.path.join(root, file)\n",
    "            \n",
    "            for attempt in range(3):\n",
    "                try:\n",
    "                    print(f\"\\033[92mProcessing file: {file_path}\\033[0m\")\n",
    "                    #prompt = \"Generate comprehensive documentation about the attached file. Make sure you include what dependencies and other files are being referenced as well as function names, class names, and what they do.\"\n",
    "                    prompt = \"Come up with a list of questions and answers about the attached file. Keep answers dense with information. A good question for a database related file would be 'What is the database technology and architecture?' or for a file that executes SQL commands 'What are the SQL commands and what do they do?' or for a file that contains a list of API endpoints 'What are the API endpoints and what do they do?'\"\n",
    "                    answer = ask_question_with_attachment(prompt, file_path)\n",
    "                    upload_prompt_answer_and_file_name(file_path, prompt, answer, repo_url)\n",
    "                    save_answers(answer, file_path, \"documentation/\")\n",
    "                    processed_files.append(file)\n",
    "                    break\n",
    "                except:\n",
    "                    print(f\"\\033[93mSkipping file: {file_path}\\033[0m\") \n",
    "                    time.sleep(15)\n",
    "            else:\n",
    "                failed_files.append(file_path)\n",
    "                \n",
    "    print(f\"Processed files: {processed_files}\")\n",
    "    print(f\"Failed files: {failed_files}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "960e8f99",
   "metadata": {},
   "source": [
    "Important: If you need to sign your ssh key before making a request, i.e. mwinit, perform that action in your terminal first and then run the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c197bd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the path to the Git repository\n",
    "# If you are using ssh to clone then uncomment the line below\n",
    "# ssh_url = '<your ssh url>'\n",
    "repo_url = \"https://github.com/aws-samples/langchain-agents.git\"\n",
    "\n",
    "# Call the function to process the repository\n",
    "process_repository(repo_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69fb9a25",
   "metadata": {},
   "source": [
    "## Using the Amazon Q Application to Answer Questions about the Repository"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08339973",
   "metadata": {},
   "source": [
    "Now that we've created the application, the index, and ingested the documentation, we can use the Amazon Q application to answer questions about the repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fdd0a89",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ask_question_about_repo(prompt):\n",
    "    answer = amazon_q.chat_sync(\n",
    "        applicationId=amazon_q_app_id,\n",
    "        userId=amazon_q_user_id,\n",
    "        userMessage=prompt\n",
    "    )\n",
    "    return answer['systemMessage']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce4a1227",
   "metadata": {},
   "outputs": [],
   "source": [
    "ask_question_about_repo(\"What is the architecture of the Bedrock Langchain solution?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b16c8ff",
   "metadata": {},
   "source": [
    "## Delete Q Application, Index, and Retriever\n",
    "\n",
    "To avoid incurring costs, you can delete the application, the index, and the retriever."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ed4fd55",
   "metadata": {},
   "outputs": [],
   "source": [
    "amazon_q.delete_retriever(applicationId=amazon_q_app_id, retrieverId=retriever_id)\n",
    "amazon_q.delete_index(applicationId=amazon_q_app_id, indexId=index_id)\n",
    "amazon_q.delete_application(applicationId=amazon_q_app_id)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}