wehub/googlecloudplatform--generative-ai
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:30:30 +08:00
commit 914fea506e
2793 changed files with 802106 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gemini/use-cases/code/analyze_codebase.ipynb
+951
View File
@@ -0,0 +1,951 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "bCIMTPB1WoTq"
},
"outputs": [],
"source": [
"# Copyright 2024 Google LLC\n",
"#\n",
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
"# https://www.apache.org/licenses/LICENSE-2.0\n",
"#\n",
"# Unless required by applicable law or agreed to in writing, software\n",
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pPmcMNgZpo9V"
},
"source": [
"# Analyze a codebase with Gemini in Vertex AI\n",
"\n",
"<table align=\"left\">\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/analyze_codebase.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fuse-cases%2Fcode%2Fanalyze_codebase.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td> \n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/use-cases/code/analyze_codebase.ipynb\">\n",
" <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/analyze_codebase.ipynb\">\n",
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://goo.gle/40yTNki\">\n",
" <img width=\"32px\" src=\"https://cdn.qwiklabs.com/assets/gcp_cloud-e3a77215f0b8bfa9b3f611c0d2208c7e8708ed31.svg\" alt=\"Google Cloud logo\"><br> Open in Cloud Skills Boost\n",
" </a>\n",
" </td>\n",
"</table>\n",
"\n",
"<div style=\"clear: both;\"></div>\n",
"\n",
"<b>Share to:</b>\n",
"\n",
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/analyze_codebase.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/analyze_codebase.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/analyze_codebase.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/analyze_codebase.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/analyze_codebase.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
"</a> "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "1EExYZvij2ve"
},
"source": [
"| | |\n",
"|-|-|\n",
"|Author(s) | [Eric Dong](https://github.com/gericdong), [Holt Skinner](https://github.com/holtskinner), [Aakash Gouda](https://github.com/aksstar)|"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7yVV6txOmNMn"
},
"source": [
"## Overview\n",
"\n",
"Gemini features a breakthrough long context window of up to 1 million tokens that can help seamlessly analyze, classify and summarize large amounts of content within a given prompt.\n",
"\n",
"With its long-context reasoning, Gemini can analyze an entire codebase for deeper insights.\n",
"\n",
"In this tutorial, you learn how to analyze an entire codebase with Gemini 3 and prompt the model to:\n",
"\n",
"- **Analyze**: Summarize codebases effortlessly.\n",
"- **Guide**: Generate clear developer getting-started documentation.\n",
"- **Debug**: Uncover critical bugs and provide fixes.\n",
"- **Enhance**: Implement new features and improve reliability and security.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hdMGtr18rFdL"
},
"source": [
"## Getting Started"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "No17Cw5hgx12"
},
"source": [
"### Install Google Gen AI SDK for Python and other libraries\n",
"\n",
"In addition to the [Google Gen AI SDK for Python](https://cloud.google.com/vertex-ai/generative-ai/docs/sdks/overview), we will be using [Gitingest](https://gitingest.com/) to load the repository into the prompt.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tFy3H3aPgx12"
},
"outputs": [],
"source": [
"%pip install --upgrade --quiet google-genai gitingest gitpython PyGithub"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "R5Xep4W9lq-Z"
},
"source": [
"### Restart runtime (Colab only)\n",
"\n",
"To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
"\n",
"The restart might take a minute or longer. After it's restarted, continue to the next step."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "XRvKdaPDTznN"
},
"outputs": [],
"source": [
"import sys\n",
"\n",
"if \"google.colab\" in sys.modules:\n",
" import IPython\n",
"\n",
" app = IPython.Application.instance()\n",
" app.kernel.do_shutdown(True)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "SbmM4z7FOBpM"
},
"source": [
"<div class=\"alert alert-block alert-warning\">\n",
"<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
"</div>\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dmWOrTJ3gx13"
},
"source": [
"### Authenticate your notebook environment (Colab only)\n",
"\n",
"If you are running this notebook on Google Colab, run the following cell to authenticate your environment.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "NyKGtVQjgx13"
},
"outputs": [],
"source": [
"import sys\n",
"\n",
"if \"google.colab\" in sys.modules:\n",
" from google.colab import auth\n",
"\n",
" auth.authenticate_user()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "DF4l8DTdWgPY"
},
"source": [
"### Set Google Cloud project information and create client\n",
"\n",
"To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
"\n",
"Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Nqwi-5ufWp_B"
},
"outputs": [],
"source": [
"import os\n",
"\n",
"# fmt: off\n",
"PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
"# fmt: on\n",
"if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
" PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
"\n",
"LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
"\n",
"from google import genai\n",
"\n",
"client = genai.Client(enterprise=True, project=PROJECT_ID, location=LOCATION)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vbozY-XKee95"
},
"source": [
"### Import libraries"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"id": "NSCFmvOWBas9"
},
"outputs": [],
"source": [
"import os\n",
"import shutil\n",
"\n",
"import git\n",
"from IPython.core.interactiveshell import InteractiveShell\n",
"from IPython.display import Markdown, display\n",
"from github import Github\n",
"from gitingest import ingest\n",
"\n",
"InteractiveShell.ast_node_interactivity = \"all\"\n",
"\n",
"import nest_asyncio\n",
"from google.genai.types import CreateCachedContentConfig, GenerateContentConfig\n",
"\n",
"nest_asyncio.apply()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "WNoOTMp2fe33"
},
"source": [
"## Cloning a codebase\n",
"\n",
"You will use the repo [Online Boutique](https://github.com/GoogleCloudPlatform/microservices-demo) as an example in this notebook.the Online Boutique is a cloud-first microservices demo application. The application is a web-based e-commerce app where users can browse items, add them to the cart, and purchase them. This application consists of 11 microservices across multiple languages."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"id": "GlDOs49qgStM"
},
"outputs": [],
"source": [
"# The GitHub repository URL\n",
"# fmt: off\n",
"repo_url = \"https://github.com/GoogleCloudPlatform/microservices-demo\" # @param {type:\"string\"}\n",
"# fmt: on\n",
"\n",
"# The location to clone the repo\n",
"repo_dir = \"./repo\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "PAm1ly9pfIEX"
},
"source": [
"#### Define helper functions for processing GitHub repository"
]
},
{
"cell_type": "code",
"execution_count": 32,
"metadata": {
"id": "stNia6UaHau2"
},
"outputs": [],
"source": [
"def clone_repo(repo_url: str, repo_dir: str) -> None:\n",
" \"\"\"Shallow clone a GitHub repository.\"\"\"\n",
" if os.path.exists(repo_dir):\n",
" shutil.rmtree(repo_dir)\n",
" os.makedirs(repo_dir, exist_ok=True)\n",
" git.Repo.clone_from(repo_url, repo_dir, depth=2)\n",
"\n",
"\n",
"def get_github_issue(owner: str, repo: str, issue_number: int) -> str | None:\n",
" \"\"\"Fetch the contents of a GitHub issue.\n",
"\n",
" Args:\n",
" owner (str): The owner of the repository.\n",
" repo (str): The name of the repository.\n",
" issue_number (int): The issue number to retrieve.\n",
"\n",
" Returns:\n",
" str | None: The issue body if found, otherwise None.\n",
"\n",
" Raises:\n",
" Exception: If an error occurs while fetching the issue.\n",
" \"\"\"\n",
" g = Github()\n",
"\n",
" try:\n",
" repository = g.get_repo(f\"{owner}/{repo}\")\n",
" issue = repository.get_issue(number=issue_number)\n",
" return issue.body\n",
" except Exception as error:\n",
" print(f\"Error fetching issue: {error}\")\n",
" return None\n",
"\n",
"\n",
"def get_git_diff(repo_dir: str) -> str:\n",
" \"\"\"Fetches commit IDs from a local Git repository on a specified branch.\"\"\"\n",
" repo = git.Repo(repo_dir)\n",
" branch_name = \"main\"\n",
"\n",
" # A list of commit IDs (SHA-1 hashes) in reverse chronological order (newest first)\n",
" commit_ids = [commit.hexsha for commit in repo.iter_commits(branch_name)]\n",
" if len(commit_ids) >= 2:\n",
" return repo.git.diff(commit_ids[0], commit_ids[1])\n",
" return \"\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2947e8614485"
},
"source": [
"### Create an index and extract content of a codebase\n",
"\n",
"Clone the repo and create an index and extract content of code/text files.\n",
"\n",
"Gitingest will extract all of the contents of the files into a long string and create a directory tree of the files."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4eb417d380c9"
},
"outputs": [],
"source": [
"clone_repo(repo_url, repo_dir)\n",
"\n",
"_, tree, content = ingest(repo_dir)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JiVQB5SKekS0"
},
"source": [
"## Analyzing the codebase with Gemini\n",
"\n",
"With its long-context reasoning, Gemini can process the codebase and answer questions about the codebase."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BY1nfXrqRxVX"
},
"source": [
"### Load the Gemini model\n",
"\n",
"Learn more about the [Gemini API models on Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-models)."
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {
"id": "vB9gY3WruzK9"
},
"outputs": [],
"source": [
"MODEL_ID = \"gemini-3.5-flash\" # @param {type:\"string\"}"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "59f3c0bbef54"
},
"source": [
"### Create a context cache\n",
"\n",
"We will create a [context cache](https://cloud.google.com/vertex-ai/generative-ai/docs/context-cache/context-cache-overview) of the codebase so we don't have to send the entire context with every request, saving processing time and cost.\n",
"\n",
"**Note**: Context caching is only available for stable models with fixed versions (for example, `gemini-3.5-flash`). You must include the version postfix (for example, the `-001`).\n",
"\n",
"For more information, see [Available Gemini stable model versions](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/model-versioning#stable-versions-available)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ee211faaf336"
},
"outputs": [],
"source": [
"system_instruction = \"You are a coding expert. Your mission is to answer all code related questions with given context and instructions.\"\n",
"\n",
"contents = [\n",
" \"\"\"\n",
" Context:\n",
" - The entire codebase is provided below.\n",
" - Here is directory tree of all of the files in the codebase:\n",
" \"\"\",\n",
" tree,\n",
" \"\"\"\n",
" - Then each of the files are concatenated together. You will find all of the code you need:\n",
" \"\"\",\n",
" content,\n",
"]\n",
"\n",
"cached_content = client.caches.create(\n",
" model=MODEL_ID,\n",
" config=CreateCachedContentConfig(\n",
" contents=contents,\n",
" system_instruction=system_instruction,\n",
" ttl=\"3600s\",\n",
" ),\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "g3OtaszvJt9L"
},
"source": [
"### 1. Summarizing the codebase\n",
"\n",
"Generate a summary of the codebase."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "uMexx1Qtf1ML"
},
"outputs": [],
"source": [
"question = \"\"\"\n",
" Give me a summary of this codebase, and tell me the top 3 things that I can learn from it.\n",
"\"\"\"\n",
"\n",
"# Generate text using non-streaming method\n",
"response = client.models.generate_content(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" # Use the cached content\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"# Print generated text and usage metadata\n",
"display(Markdown(response.text))\n",
"print(response.usage_metadata)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aCshJHPCYoxI"
},
"source": [
"### 2. Creating a developer getting started guide\n",
"\n",
"Generate a getting started guide for developers. This sample uses the streaming option to generate the content."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "e6Kns7vCYm1P"
},
"outputs": [],
"source": [
"question = \"\"\"\n",
" Provide a getting started guide to onboard new developers to the codebase.\n",
"\"\"\"\n",
"\n",
"\n",
"# Generate text using streaming method\n",
"responses = client.models.generate_content_stream(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"for response in responses:\n",
" print(response.text, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hXurINu-jelb"
},
"source": [
"### 3. Finding bugs\n",
"\n",
"Find the top 3 most severe issues in the codebase."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "fy3AWPRgNhu_"
},
"outputs": [],
"source": [
"question = \"\"\"\n",
" Find the top 3 most severe issues in the codebase.\n",
"\"\"\"\n",
"\n",
"responses = client.models.generate_content_stream(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"for response in responses:\n",
" print(response.text, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HCilrR6FjmfB"
},
"source": [
"### 4. Fixing bug\n",
"\n",
"Find the most severe issue in the codebase that can be fixed and provide a code fix for it.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "dwjDh0xGKE2r"
},
"outputs": [],
"source": [
"question = \"\"\"\n",
" Find the most severe bug in the codebase that you can provide a code fix for.\n",
"\"\"\"\n",
"\n",
"responses = client.models.generate_content_stream(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"for response in responses:\n",
" print(response.text, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8w2pCULT_xKE"
},
"source": [
"### 5. Implementing a feature request using Function Calling\n",
"\n",
"Get the feature request text from a GitHub Issue URL.\n",
"\n",
"We will use [Function Calling](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/function-calling) to extract the feature request data from the prompt, then call the GitHub API to retrieve the contents.\n",
"\n",
"Note: We can't use the previously created cached content, because tools cannot be added at runtime when using cached content, and the other prompts in this notebook do not need this tool."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "mMjOy0gJ1_xx"
},
"outputs": [],
"source": [
"FEATURE_REQUEST_URL = (\n",
" \"https://github.com/GoogleCloudPlatform/microservices-demo/issues/2205\"\n",
")\n",
"\n",
"question = f\"What is the feature request of the following {FEATURE_REQUEST_URL}\"\n",
"\n",
"response = client.models.generate_content(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" # Use the function as a tool\n",
" config=GenerateContentConfig(\n",
" tools=[get_github_issue],\n",
" ),\n",
")\n",
"\n",
"issue_description = response.text\n",
"display(Markdown(f\"# Feature Request\\n{issue_description}\"))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "19455545f4c9"
},
"source": [
"Use the GitHub Issue text to implement the feature request."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "78e6df259be8"
},
"outputs": [],
"source": [
"# Combine feature request content and cached code content\n",
"question = f\"\"\"Implement the following feature request\n",
"{issue_description}\n",
"\"\"\"\n",
"\n",
"response = client.models.generate_content(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"# Generate code response\n",
"display(Markdown(response.text))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "TOk_Qe35b_cJ"
},
"source": [
"### 6. Creating a troubleshooting guide\n",
"\n",
"Create a troubleshooting guide to help resolve common issues."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "DKn85LS-v0iw"
},
"outputs": [],
"source": [
"question = \"\"\"\n",
" Provide a troubleshooting guide to help resolve common issues.\n",
"\"\"\"\n",
"\n",
"responses = client.models.generate_content_stream(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"for response in responses:\n",
" print(response.text, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "h23z0sTsj5pL"
},
"source": [
"### 7. Making the app more reliable\n",
"\n",
"Recommend best practices to make the application more reliable.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yOBSulTPLUAo"
},
"outputs": [],
"source": [
"question = \"\"\"\n",
" How can I make this application more reliable? Consider best practices from https://www.r9y.dev/\n",
"\"\"\"\n",
"\n",
"responses = client.models.generate_content_stream(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"for response in responses:\n",
" print(response.text, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uf1jNDpJj8u0"
},
"source": [
"### 8. Making the app more secure\n",
"\n",
"Recommend best practices to make the application more secure."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Hy_mCyFVLlXU"
},
"outputs": [],
"source": [
"question = \"\"\"\n",
" How can you secure the application?\n",
"\"\"\"\n",
"\n",
"responses = client.models.generate_content_stream(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"for response in responses:\n",
" print(response.text, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GFfwMOb6kYfw"
},
"source": [
"### 9. Learning the codebase\n",
"\n",
"Create a quiz about the concepts used in the codebase."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "l7jQIUwsNRH4"
},
"outputs": [],
"source": [
"question = \"\"\"\n",
" Create a quiz about the concepts used in the codebase to help me solidify my understanding.\n",
"\"\"\"\n",
"\n",
"responses = client.models.generate_content_stream(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"for response in responses:\n",
" print(response.text, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Rjo1UrZwGLan"
},
"source": [
"### 10. Creating a quickstart tutorial\n",
"\n",
"Create an end-to-end quickstart tutorial for a specific component.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "FRwmRyDDFRMB"
},
"outputs": [],
"source": [
"question = \"\"\"\n",
" Please write an end-to-end quickstart tutorial that introduces AlloyDB,\n",
" shows how to configure it with the CartService,\n",
" and highlights key capabilities of AlloyDB in context of the Online Boutique application.\n",
"\"\"\"\n",
"\n",
"responses = client.models.generate_content_stream(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"for response in responses:\n",
" print(response.text, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "OAJ-kBCZnlH_"
},
"source": [
"### 11. Creating a Git Changelog Generator\n",
"\n",
"Understanding changes made between Git commits and highlighting the most important aspects of the changes."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "jdeWO8crnlH_"
},
"outputs": [],
"source": [
"diff_text = get_git_diff(repo_dir)\n",
"question = f\"\"\"\n",
" Given the below git diff output, Summarize the important changes made.\n",
"```diff\n",
"{diff_text}\n",
"```\n",
"\"\"\"\n",
"\n",
"responses = client.models.generate_content_stream(\n",
" model=MODEL_ID,\n",
" contents=question,\n",
" config=GenerateContentConfig(\n",
" cached_content=cached_content.name,\n",
" ),\n",
")\n",
"\n",
"for response in responses:\n",
" print(response.text, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7kUeIBfGyoX7"
},
"source": [
"## Conclusion\n",
"\n",
"In this tutorial, you've learned how to use Gemini to analyze a codebase and prompt the model to:\n",
"\n",
"- Summarize codebases effortlessly.\n",
"- Generate clear developer getting-started documentation.\n",
"- Uncover critical bugs and provide fixes.\n",
"- Implement new features and improve reliability and security.\n",
"- Understanding changes made between Git commits"
]
}
],
"metadata": {
"colab": {
"name": "analyze_codebase.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
gemini/use-cases/code/code_retrieval_augmented_generation.ipynb
+780
View File
@@ -0,0 +1,780 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5norOZI0mA6s"
},
"outputs": [],
"source": [
"# Copyright 2023 Google LLC\n",
"#\n",
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
"# https://www.apache.org/licenses/LICENSE-2.0\n",
"#\n",
"# Unless required by applicable law or agreed to in writing, software\n",
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "XNPE46X8mJj4"
},
"source": [
"# Use Retrieval Augmented Generation (RAG) with Gemini API\n",
"\n",
"<table align=\"left\">\n",
"\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_retrieval_augmented_generation.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
"\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fuse-cases%2Fcode%2Fcode_retrieval_augmented_generation.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_retrieval_augmented_generation.ipynb\">\n",
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/use-cases/code/code_retrieval_augmented_generation.ipynb\">\n",
" <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
" </a>\n",
" </td>\n",
"</table>\n",
"\n",
"<div style=\"clear: both;\"></div>\n",
"\n",
"<b>Share to:</b>\n",
"\n",
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_retrieval_augmented_generation.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_retrieval_augmented_generation.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_retrieval_augmented_generation.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_retrieval_augmented_generation.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_retrieval_augmented_generation.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
"</a> "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VrLtlKPFqSxB"
},
"source": [
"| | |\n",
"|-|-|\n",
"|Author(s) | [Lavi Nigam](https://github.com/lavinigam-gcp), [Polong Lin](https://github.com/polong-lin) |"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "zNAEdYNFmQcP"
},
"source": [
"### Objective\n",
"\n",
"This notebook demonstrates how you augment output from Gemini API by bringing in external knowledge. An example is provided using Code Retrieval Augmented Generation(RAG) pattern using [Google Cloud's Generative AI github repository](https://github.com/GoogleCloudPlatform/generative-ai) as external knowledge. The notebook uses [Gemini API in Vertex AI](https://ai.google.dev/gemini-api), [Embeddings for Text API](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings), FAISS vector store and [LangChain 🦜️🔗](https://python.langchain.com/en/latest/).\n",
"\n",
"### Overview\n",
"\n",
"Here is overview of what we'll go over.\n",
"\n",
"Index Creation:\n",
"\n",
"1. Recursively list the files(.ipynb) in github repo\n",
"2. Extract code and markdown from the files\n",
"3. Chunk & generate embeddings for each code strings and add initialize the vector store\n",
"\n",
"Runtime:\n",
"\n",
"4. User enters a prompt or asks a question as a prompt\n",
"5. Try zero-shot prompt\n",
"6. Run prompt using RAG Chain & compare results.To generate response we use **gemini-2.0-flash**\n",
"\n",
"### Cost\n",
"\n",
"This tutorial uses billable components of Google Cloud:\n",
"\n",
"- Gemini API in Vertex AI offered by Google Cloud\n",
"\n",
"Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n",
"\n",
"**Note:** We are using local vector store(FAISS) for this example however recommend managed highly scalable vector store for production usage such as [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview) or [AlloyDB for PostgreSQL](https://cloud.google.com/alloydb/docs/ai/work-with-embeddings) or [Cloud SQL for PostgreSQL](https://cloud.google.com/sql/docs/postgres/features) using pgvector extension."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2cab0c8509c9"
},
"source": [
"## Get started"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "b56b5a5d28c1"
},
"source": [
"### Install Vertex AI SDK for Python and other required packages\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "QHaqV20Csqkt"
},
"outputs": [],
"source": [
"%pip install --upgrade --user -q google-cloud-aiplatform \\\n",
" langchain \\\n",
" langchain_google_vertexai \\\n",
" langchain-community \\\n",
" faiss-cpu \\\n",
" nbformat"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-VUWOgz6M1rZ"
},
"source": [
"### Restart runtime (Colab only)\n",
"\n",
"To use the newly installed packages, you must restart the runtime on Google Colab."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "BIS8EYgkMy8T"
},
"outputs": [],
"source": [
"import sys\n",
"\n",
"if \"google.colab\" in sys.modules:\n",
" import IPython\n",
"\n",
" app = IPython.Application.instance()\n",
" app.kernel.do_shutdown(True)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0af13c10a26a"
},
"source": [
"<div class=\"alert alert-block alert-warning\">\n",
"<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
"</div>\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uZcP9WBENG0e"
},
"source": [
"### Authenticate your notebook environment (Colab only)\n",
"\n",
"Authenticate your environment on Google Colab.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1S_HgQXQNcbz"
},
"outputs": [],
"source": [
"import sys\n",
"\n",
"if \"google.colab\" in sys.modules:\n",
" from google.colab import auth\n",
"\n",
" auth.authenticate_user()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rVmxMr43Nhoo"
},
"source": [
"### Import libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "L-Tljm5asMBc"
},
"outputs": [],
"source": [
"import time\n",
"\n",
"from google.cloud import aiplatform\n",
"from langchain.chains import RetrievalQA\n",
"from langchain.prompts import PromptTemplate\n",
"from langchain.schema.document import Document\n",
"from langchain.text_splitter import Language, RecursiveCharacterTextSplitter\n",
"from langchain.vectorstores import FAISS\n",
"\n",
"# LangChain\n",
"from langchain_google_vertexai import VertexAI, VertexAIEmbeddings\n",
"import nbformat\n",
"import requests\n",
"\n",
"# Vertex AI\n",
"import vertexai\n",
"\n",
"# Print the version of Vertex AI SDK for Python\n",
"print(f\"Vertex AI SDK version: {aiplatform.__version__}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "4f872cd812d0"
},
"source": [
"### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
"\n",
"To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "eNGEcBKG0iK-"
},
"outputs": [],
"source": [
"# Initialize project\n",
"# Define project information\n",
"PROJECT_ID = \"YOUR_PROJECT_ID\" # @param {type:\"string\"}\n",
"LOCATION = \"us-central1\" # @param {type:\"string\"}\n",
"\n",
"vertexai.init(project=PROJECT_ID, location=LOCATION)\n",
"\n",
"# Code Generation\n",
"code_llm = VertexAI(\n",
" model_name=\"gemini-2.0-flash\",\n",
" max_output_tokens=2048,\n",
" temperature=0.1,\n",
" verbose=False,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "o537exyZk9DI"
},
"source": [
"Next we need to create a GitHub personal token to be able to list all files in a repository.\n",
"\n",
"- Follow [this link](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens) to create GitHub token with repo->public_repo scope and update `GITHUB_TOKEN` variable below."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Bt9IVDSqk7y4"
},
"outputs": [],
"source": [
"# provide GitHub personal access token\n",
"GITHUB_TOKEN = \"YOUR_GITHUB_TOKEN\" # @param {type:\"string\"}\n",
"GITHUB_REPO = \"GoogleCloudPlatform/generative-ai\" # @param {type:\"string\"}"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dqq3GeEbOJbU"
},
"source": [
"# Index Creation\n",
"\n",
"We use the Google Cloud Generative AI github repository as the data source. First list all Jupyter Notebook files in the repo and store it in a text file.\n",
"\n",
"You can skip this step(#1) if you have executed it once and generated the output text file.\n",
"\n",
"### 1. Recursively list the files(.ipynb) in the github repository"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "eTA1Jt0uOX8y"
},
"outputs": [],
"source": [
"# Crawls a GitHub repository and returns a list of all ipynb files in the repository\n",
"\n",
"\n",
"def crawl_github_repo(url: str, is_sub_dir: bool, access_token: str = GITHUB_TOKEN):\n",
" ignore_list = [\"__init__.py\"]\n",
"\n",
" if not is_sub_dir:\n",
" api_url = f\"https://api.github.com/repos/{url}/contents\"\n",
"\n",
" else:\n",
" api_url = url\n",
"\n",
" headers = {\n",
" \"Accept\": \"application/vnd.github.v3+json\",\n",
" \"Authorization\": f\"Bearer {access_token}\",\n",
" }\n",
"\n",
" response = requests.get(api_url, headers=headers)\n",
" response.raise_for_status() # Check for any request errors\n",
"\n",
" files = []\n",
"\n",
" contents = response.json()\n",
"\n",
" for item in contents:\n",
" if (\n",
" item[\"type\"] == \"file\"\n",
" and item[\"name\"] not in ignore_list\n",
" and (item[\"name\"].endswith(\".py\") or item[\"name\"].endswith(\".ipynb\"))\n",
" ):\n",
" files.append(item[\"html_url\"])\n",
" elif item[\"type\"] == \"dir\" and not item[\"name\"].startswith(\".\"):\n",
" sub_files = crawl_github_repo(item[\"url\"], True)\n",
" time.sleep(0.1)\n",
" files.extend(sub_files)\n",
"\n",
" return files"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5vaKaxcGO_R6"
},
"outputs": [],
"source": [
"code_files_urls = crawl_github_repo(GITHUB_REPO, False, GITHUB_TOKEN)\n",
"\n",
"# Write list to a file so you do not have to download each time\n",
"with open(\"code_files_urls.txt\", \"w\") as f:\n",
" for item in code_files_urls:\n",
" f.write(item + \"\\n\")\n",
"\n",
"len(code_files_urls)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "c5hoNYJ5byMJ"
},
"outputs": [],
"source": [
"code_files_urls[0:10]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "mFNVieLnR8Ie"
},
"source": [
"### 2. Extract code from the Jupyter notebooks.\n",
"\n",
"You could also include .py file, shell scripts etc."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ZsM1M4hn4cBu"
},
"outputs": [],
"source": [
"# Extracts the python code from an ipynb file from github\n",
"\n",
"\n",
"def extract_python_code_from_ipynb(github_url, cell_type=\"code\"):\n",
" raw_url = github_url.replace(\"github.com\", \"raw.githubusercontent.com\").replace(\n",
" \"/blob/\", \"/\"\n",
" )\n",
"\n",
" response = requests.get(raw_url)\n",
" response.raise_for_status() # Check for any request errors\n",
"\n",
" notebook_content = response.text\n",
"\n",
" notebook = nbformat.reads(notebook_content, as_version=nbformat.NO_CONVERT)\n",
"\n",
" python_code = None\n",
"\n",
" for cell in notebook.cells:\n",
" if cell.cell_type == cell_type:\n",
" if not python_code:\n",
" python_code = cell.source\n",
" else:\n",
" python_code += \"\\n\" + cell.source\n",
"\n",
" return python_code\n",
"\n",
"\n",
"def extract_python_code_from_py(github_url):\n",
" raw_url = github_url.replace(\"github.com\", \"raw.githubusercontent.com\").replace(\n",
" \"/blob/\", \"/\"\n",
" )\n",
"\n",
" response = requests.get(raw_url)\n",
" response.raise_for_status() # Check for any request errors\n",
"\n",
" python_code = response.text\n",
"\n",
" return python_code"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WCRp5Xtb48is"
},
"outputs": [],
"source": [
"with open(\"code_files_urls.txt\") as f:\n",
" code_files_urls = f.read().splitlines()\n",
"len(code_files_urls)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4Y9SMO7H4xgF"
},
"outputs": [],
"source": [
"code_strings = []\n",
"\n",
"for i in range(0, len(code_files_urls)):\n",
" if code_files_urls[i].endswith(\".ipynb\"):\n",
" content = extract_python_code_from_ipynb(code_files_urls[i], \"code\")\n",
" doc = Document(\n",
" page_content=content, metadata={\"url\": code_files_urls[i], \"file_index\": i}\n",
" )\n",
" code_strings.append(doc)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "T1AF3fhBSLOm"
},
"source": [
"### 3. Chunk & generate embeddings for each code strings & initialize the vector store\n",
"\n",
"We need to split code into usable chunks that the LLM can use for code generation. Therefore it's crucial to use the right chunking approach and chunk size."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Rj1cCA2fqx64"
},
"outputs": [],
"source": [
"# Utility functions for Embeddings API with rate limiting\n",
"\n",
"\n",
"def rate_limit(max_per_minute):\n",
" period = 60 / max_per_minute\n",
" print(\"Waiting\")\n",
" while True:\n",
" before = time.time()\n",
" yield\n",
" after = time.time()\n",
" elapsed = after - before\n",
" sleep_time = max(0, period - elapsed)\n",
" if sleep_time > 0:\n",
" print(\".\", end=\"\")\n",
" time.sleep(sleep_time)\n",
"\n",
"\n",
"class CustomVertexAIEmbeddings(VertexAIEmbeddings):\n",
" requests_per_minute: int\n",
" num_instances_per_batch: int\n",
" model_name: str\n",
"\n",
" # Overriding embed_documents method\n",
" def embed_documents(\n",
" self, texts: list[str], batch_size: int | None = None\n",
" ) -> list[list[float]]:\n",
" limiter = rate_limit(self.requests_per_minute)\n",
" results = []\n",
" docs = list(texts)\n",
"\n",
" while docs:\n",
" # Working in batches because the API accepts maximum 5\n",
" # documents per request to get embeddings\n",
" head, docs = (\n",
" docs[: self.num_instances_per_batch],\n",
" docs[self.num_instances_per_batch :],\n",
" )\n",
" chunk = self.client.get_embeddings(head)\n",
" results.extend(chunk)\n",
" next(limiter)\n",
"\n",
" return [r.values for r in results]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "oae37l-pvzZ6"
},
"outputs": [],
"source": [
"# Chunk code strings\n",
"text_splitter = RecursiveCharacterTextSplitter.from_language(\n",
" language=Language.PYTHON, chunk_size=2000, chunk_overlap=200\n",
")\n",
"\n",
"\n",
"texts = text_splitter.split_documents(code_strings)\n",
"print(len(texts))\n",
"\n",
"# Initialize Embedding API\n",
"EMBEDDING_QPM = 100\n",
"EMBEDDING_NUM_BATCH = 5\n",
"embeddings = CustomVertexAIEmbeddings(\n",
" requests_per_minute=EMBEDDING_QPM,\n",
" num_instances_per_batch=EMBEDDING_NUM_BATCH,\n",
" model_name=\"text-embedding-005\",\n",
")\n",
"\n",
"# Create Index from embedded code chunks\n",
"db = FAISS.from_documents(texts, embeddings)\n",
"\n",
"# Init your retriever.\n",
"retriever = db.as_retriever(\n",
" search_type=\"similarity\", # Also test \"similarity\", \"mmr\"\n",
" search_kwargs={\"k\": 5},\n",
")\n",
"\n",
"retriever"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Q_gn89IyuHIT"
},
"source": [
"# Runtime\n",
"### 4. User enters a prompt or asks a question as a prompt"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1vrvTkO7uFNi"
},
"outputs": [],
"source": [
"user_question = \"Create a Python function that takes a prompt and predicts using langchain.llms interface with Vertex AI Gemini model\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "azbvOUFRvEp5"
},
"outputs": [],
"source": [
"# Define prompt templates\n",
"\n",
"# Zero Shot prompt template\n",
"prompt_zero_shot = \"\"\"\n",
" You are a proficient python developer. Respond with the syntactically correct & concise code for to the question below.\n",
"\n",
" Question:\n",
" {question}\n",
"\n",
" Output Code :\n",
" \"\"\"\n",
"\n",
"prompt_prompt_zero_shot = PromptTemplate(\n",
" input_variables=[\"question\"],\n",
" template=prompt_zero_shot,\n",
")\n",
"\n",
"\n",
"# RAG template\n",
"prompt_RAG = \"\"\"\n",
" You are a proficient python developer. Respond with the syntactically correct code for to the question below. Make sure you follow these rules:\n",
" 1. Use context to understand the APIs and how to use it & apply.\n",
" 2. Do not add license information to the output code.\n",
" 3. Do not include Colab code in the output.\n",
" 4. Ensure all the requirements in the question are met.\n",
"\n",
" Question:\n",
" {question}\n",
"\n",
" Context:\n",
" {context}\n",
"\n",
" Helpful Response :\n",
" \"\"\"\n",
"\n",
"prompt_RAG_template = PromptTemplate(\n",
" template=prompt_RAG, input_variables=[\"context\", \"question\"]\n",
")\n",
"\n",
"qa_chain = RetrievalQA.from_llm(\n",
" llm=code_llm,\n",
" prompt=prompt_RAG_template,\n",
" retriever=retriever,\n",
" return_source_documents=True,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3NBaObAQSlIv"
},
"source": [
"### 5. Try zero-shot prompt"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1svTVwtBS0zP"
},
"outputs": [],
"source": [
"response = code_llm.invoke(input=user_question, max_output_tokens=2048, temperature=0.1)\n",
"print(response)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JPm8qdxzwPM0"
},
"source": [
"### 6. Run prompt using RAG Chain & compare results\n",
"To generate response we use code-bison however can also use code-gecko and codechat-bison"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ZMz3nPMyVoj_"
},
"outputs": [],
"source": [
"results = qa_chain.invoke(input={\"query\": user_question})\n",
"print(results[\"result\"])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HF3lVWK1wjxe"
},
"source": [
"### Let's try another prompt"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "jel0ON68XiU7"
},
"outputs": [],
"source": [
"user_question = \"Create python function that takes text input and returns embeddings using LangChain with Vertex AI text-embedding-005 model\"\n",
"\n",
"\n",
"response = code_llm.invoke(input=user_question, max_output_tokens=2048, temperature=0.1)\n",
"print(response)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "G9bIkqE8sO6P"
},
"outputs": [],
"source": [
"results = qa_chain.invoke(input={\"query\": user_question})\n",
"print(results[\"result\"])"
]
}
],
"metadata": {
"colab": {
"name": "code_retrieval_augmented_generation.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb
+901
View File
@@ -0,0 +1,901 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "nUWIDIzvODfw",
"metadata": {
"id": "nUWIDIzvODfw"
},
"source": [
"# Code Vulnerability Scanning & Automated Remediation using Gemini API in Vertex AI (Gemini 2.0)\n",
"\n",
"---\n"
]
},
{
"cell_type": "markdown",
"id": "xHjHLyn8m2WF",
"metadata": {
"id": "xHjHLyn8m2WF"
},
"source": [
"<table align=\"left\">\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\">\n",
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fuse-cases%2Fcode%2Fcode_scanning_and_vulnerability_detection.ipynb\">\n",
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
" </a>\n",
" </td> \n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\">\n",
" <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\">\n",
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://goo.gle/4j91rJr\">\n",
" <img width=\"32px\" src=\"https://cdn.qwiklabs.com/assets/gcp_cloud-e3a77215f0b8bfa9b3f611c0d2208c7e8708ed31.svg\" alt=\"Google Cloud logo\"><br> Open in Cloud Skills Boost\n",
" </a>\n",
" </td>\n",
"</table>\n",
"\n",
"<div style=\"clear: both;\"></div>\n",
"\n",
"<b>Share to:</b>\n",
"\n",
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
"</a>\n",
"\n",
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\" target=\"_blank\">\n",
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
"</a> \n"
]
},
{
"cell_type": "markdown",
"id": "zBrMoWSrXFI6",
"metadata": {
"id": "zBrMoWSrXFI6"
},
"source": [
"| | |\n",
"|-|-|\n",
"|Author(s) | [Souvik Mukherjee](https://github.com/talktosauvik)"
]
},
{
"cell_type": "markdown",
"id": "R2__rS9iZeVx",
"metadata": {
"id": "R2__rS9iZeVx"
},
"source": [
"## Background\n",
"\n",
"In today's digital landscape, software security is paramount. With the increasing sophistication of cyber threats, it's more important than ever for developers to proactively identify and address vulnerabilities in their code. Vulnerabilities can lead to data breaches, financial losses, and reputational damage. By harnessing the power of Gemini 2.0, we can help transform code vulnerability detection and remediation, and build a software vulnerability scanning mechanism"
]
},
{
"cell_type": "markdown",
"id": "6j_ed1PcwS3f",
"metadata": {
"id": "6j_ed1PcwS3f"
},
"source": [
"## Overview\n",
"\n",
"Gemini 2.0, a member of Google Gemini family, is a generative AI model purpose-built for diverse multimodal applications. It's proficiency in understanding and generating content across text, code, and images makes it a powerful asset for intricate codebase analysis. With its expansive 2M token context window, Gemini 2.0 efficiently processes large code volumes in a single call, streamlining large-scale code scanning.Gemini 2.0's deep comprehension of programming languages and security best practices enables it to identify potential vulnerabilities and suggest helpful and contextual modifications. Learn more about [Gemini 2.0](https://deepmind.google/technologies/gemini/pro/).\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"This experimental approach aims to efficiently scan large codebases, analyze multiple files in a single call, and delve deeper into complex code relationships and patternsThe model's deep analysis of code can help ensure comprehensive vulnerability detection, going beyond surface-level flaws. By using this approach, we can accommodate code written in several programming languages. Additionally, we can generate the findings and recommendations as JSON or CSV reports, which we would hypothetically use to make comparisons against established benchmarks and policy checks. With this tutorial, you learn how to use the Gemini API in Vertex AI, Google Cloud Storage API and the Vertex AI SDK to work with the Gemini 2.0 model to build a step by step code vulnerability scanning approach using Gemini 2.0 with the following steps:\n",
"\n",
"\n",
"* Read Python files from a GCS bucket and combining them into a single string\n",
"* Prompt engineering by crafting a clear and comprehensive prompt for Gemini 2.0, providing instructions for code analysis and output formatting\n",
"\n",
"* Submit the consolidated code string to Gemini 2.0 for analysis\n",
"* Extract vulnerability information, recommendations, and code snippets from the model response\n",
"\n",
"* Generate CSV and JSON output reports for further analysis, benchmarking and integration with security tools\n"
]
},
{
"cell_type": "markdown",
"id": "TpJD9T2_PvPp",
"metadata": {
"id": "TpJD9T2_PvPp"
},
"source": [
"### Getting Started"
]
},
{
"cell_type": "markdown",
"id": "k9uwJYsyP7rv",
"metadata": {
"id": "k9uwJYsyP7rv"
},
"source": [
"### Install Vertex AI and Google Cloud Storage SDKs for Python\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "N75LKYKXbuHp5cpzkJiq89Sb",
"metadata": {
"id": "N75LKYKXbuHp5cpzkJiq89Sb"
},
"outputs": [],
"source": [
"%pip install -q --upgrade --user google-cloud-aiplatform google-cloud-storage"
]
},
{
"cell_type": "markdown",
"id": "Xb3No6mSQFoe",
"metadata": {
"id": "Xb3No6mSQFoe"
},
"source": [
"### Restart the Kernel\n",
"\n",
"To use the newly installed packages, you must restart the current runtime."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2l4Q9oy8f2KR",
"metadata": {
"id": "2l4Q9oy8f2KR"
},
"outputs": [],
"source": [
"# restart the current runtime to be able to access the downloaded packages\n",
"import IPython\n",
"\n",
"app = IPython.Application.instance()\n",
"app.kernel.do_shutdown(True)"
]
},
{
"cell_type": "markdown",
"id": "IIyao4PXQkAm",
"metadata": {
"id": "IIyao4PXQkAm"
},
"source": [
"### Authenticate the notebook environment (Colab only)\n",
"\n",
"If you are running this notebook on Google Colab, run the following cell to authenticate your environment. This step is not required if you are using [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "pX9rkTlff8l-",
"metadata": {
"id": "pX9rkTlff8l-"
},
"outputs": [],
"source": [
"from google.colab import auth\n",
"\n",
"auth.authenticate_user()"
]
},
{
"cell_type": "markdown",
"id": "x4A6vhaKQqKT",
"metadata": {
"id": "x4A6vhaKQqKT"
},
"source": [
"### Set Google Cloud project information and initialize Vertex AI SDK\n",
"\n",
"\n",
"To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
"\n",
"Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "-c8ZjhqIgA7K",
"metadata": {
"id": "-c8ZjhqIgA7K"
},
"outputs": [],
"source": [
"# initialize variables\n",
"\n",
"PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n",
"REGION = \"us-central1\" # @param {type:\"string\"}\n",
"BUCKET_NAME = \"your-bucket-name\" # @param {type:\"string\"}\n",
"PREFIX = \"your prefix folder-name/\" # @param {type:\"string\"}"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "TwEE78OKgTRB",
"metadata": {
"id": "TwEE78OKgTRB"
},
"outputs": [],
"source": [
"# import and initialize Vertex AI\n",
"import vertexai\n",
"\n",
"vertexai.init(project=PROJECT_ID, location=REGION)"
]
},
{
"cell_type": "markdown",
"id": "zrgeiLeZRB8m",
"metadata": {
"id": "zrgeiLeZRB8m"
},
"source": [
"## Process Python files in batch\n",
"\n",
"This block of code reads Python files from the GCS bucket, combines their content, and adds respective `filename` as separator for LLM to better identify each file.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "lZNJhxiignA5",
"metadata": {
"id": "lZNJhxiignA5"
},
"outputs": [],
"source": [
"from google.cloud import storage\n",
"\n",
"\n",
"def process_py_files(bucket_name, prefix):\n",
" \"\"\"\n",
" Reads .py files from a GCS bucket and combines them into a single string.\n",
" Returns:\n",
" A string containing the combined content of all .py files.\n",
" \"\"\"\n",
"\n",
" storage_client = storage.Client()\n",
" bucket = storage_client.get_bucket(BUCKET_NAME)\n",
" blobs = bucket.list_blobs(prefix=PREFIX)\n",
"\n",
" combined_text = \"\"\n",
" for blob in blobs:\n",
" if blob.name.endswith(\".py\"):\n",
" file_content = blob.download_as_string().decode(\"utf-8\")\n",
" combined_text += f\"### File: {blob.name} ###\\n{file_content}\\n\"\n",
"\n",
" return combined_text\n",
"\n",
"\n",
"combined_string = process_py_files(BUCKET_NAME, PREFIX)\n",
"print(combined_string)"
]
},
{
"cell_type": "markdown",
"id": "ODdfYDIdRzzX",
"metadata": {
"id": "ODdfYDIdRzzX"
},
"source": [
"### Import Generative model library from Vertex AI"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "WCIOjVZYh28Q",
"metadata": {
"id": "WCIOjVZYh28Q"
},
"outputs": [],
"source": [
"from IPython.display import display\n",
"from vertexai.generative_models import GenerationConfig, GenerativeModel"
]
},
{
"cell_type": "markdown",
"id": "zyh4WLTZR42A",
"metadata": {
"id": "zyh4WLTZR42A"
},
"source": [
"### Initiate Gemini 2.0"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "sjA145qSh9O8",
"metadata": {
"id": "sjA145qSh9O8"
},
"outputs": [],
"source": [
"model = GenerativeModel(\"gemini-2.0-flash\")"
]
},
{
"cell_type": "markdown",
"id": "QHjEFzCBR_n9",
"metadata": {
"id": "QHjEFzCBR_n9"
},
"source": [
"### Setting up model configuration & Prompt template\n",
"\n",
"This piece of code sets up the model configurations & prompt template with 1 shot inference. For this specific notebook, the safety filters have not been imported.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "rX2dng4piDHd",
"metadata": {
"id": "rX2dng4piDHd"
},
"outputs": [],
"source": [
"# define the prompt template to be passed to gemini.\n",
"context = combined_string\n",
"\n",
"my_prompt = f\"\"\"You are an expert code assistant. Review the following code for vulnerabilities and provide recommendations:\n",
"{context}\n",
"\n",
"Please format your response using markdown and Display with the following structure:\n",
"\n",
"file_name: Name of the code file\n",
"Vulnerability Name: Name of the identified vulnerability\n",
"Vulnerability: Description of the vulnerability and its potential impact.\n",
"Recommendations: List of actionable steps to mitigate the vulnerability.\n",
"Recommended code: Recommended code snippet to remove the suspected vulnerability\n",
"\n",
"I am also providing a sample response output that you should follow-\n",
"-------------------------------------------------\n",
"**file_name: bulk/example_1.py**\n",
"\n",
"**Vulnerability Name:** Information Exposure\n",
"\n",
"**Vulnerability:** The `server_bad` function returns the entire traceback when an exception occurs\n",
".This can leak sensitive information about the application's internals, such as file paths, variable names, and even the type of exception raised. Attackers could\n",
" exploit this information to gain a deeper understanding of the system and potentially launch further attacks.\n",
"\n",
"**Recommendations:**\n",
"\n",
"*Catch specific exceptions instead of using a broad `Exception` clause.\n",
"*Return a generic error message to the user without revealing internal details.\n",
"*Log the full traceback for debugging purposes, but do not expose it to the user.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"sample code\n",
"```\n",
"\"\"\" # try your own prompt\n",
"\n",
"\n",
"generation_config = GenerationConfig(\n",
" temperature=0.5,\n",
" top_p=0.4,\n",
" top_k=24,\n",
" candidate_count=1,\n",
" max_output_tokens=8192,\n",
")\n",
"responses = model.generate_content(\n",
" contents=my_prompt,\n",
" generation_config=generation_config,\n",
" stream=True,\n",
")\n",
"\n",
"for res in responses:\n",
" print(res.text)"
]
},
{
"cell_type": "markdown",
"id": "cnuNp-heSUMD",
"metadata": {
"id": "cnuNp-heSUMD"
},
"source": [
"## Capture the model response into a single variable\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "M_gJqJWAhALS",
"metadata": {
"id": "M_gJqJWAhALS"
},
"outputs": [],
"source": [
"response_text = r\"\"\"\n",
" Security Analysis of Code Snippets\n",
"\n",
"### file_name: Bulk/example_1.py\n",
"\n",
"**Vulnerability Name:** Information Exposure Through Exception Handling\n",
"\n",
"**Vulnerability:** The `server_bad` function returns the entire traceback when an exception occurs. This can leak sensitive information about the application's internals, such\n",
" as file paths, variable names, and even the type of exception raised. Attackers could exploit this information to gain a deeper understanding of the system and potentially launch\n",
" further attacks.\n",
"\n",
"**Recommendations:**\n",
"\n",
"*Catch specific exceptions instead of using a broad `Exception` clause.\n",
"*Return a generic error message to the user without revealing internal details.\n",
"*Log the full traceback for debugging purposes, but do not expose it to the user.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"@app.route('/bad')\n",
"def server_bad():\n",
" try:\n",
" result = do_computation()\n",
" return\n",
" result\n",
" except SpecificException as e:\n",
" # Log the error for debugging\n",
" app.logger.exception(\"An error occurred during computation\")\n",
" return \"An error occurred. Please try again later.\"\n",
"```\n",
"\n",
"### file_name: Bulk/example_10.py\n",
"\n",
"**Vulnerability Name:** Insecure Temporary File Creation\n",
"\n",
"**Vulnerability:** The `mktemp()` function creates temporary files with predictable names, which could allow attackers to guess the file names and access sensitive data. Additionally, the file permissions might be insecure, allowing unauthorized access.\n",
"\n",
"**Recommendations:**\n",
"\n",
"*Use the `tempfile .mkstemp()` function to create temporary files with more secure permissions and unpredictable names.\n",
"*Ensure that the temporary files are deleted after they are no longer needed.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"import tempfile\n",
"\n",
"def write_results(results):\n",
" with tempfile.NamedTemporaryFile(mode=\"\n",
"w+\", delete=False) as f:\n",
" filename = f.name\n",
" f.write(results)\n",
" print(\"Results written to\", filename)\n",
" # ... process the file ...\n",
" os.remove(filename) # Delete the file when done\n",
"```\n",
"\n",
"### file_name: Bulk/example_100.py\n",
"\n",
"**Vulnerability Name:** Weak Key Generation Algorithm\n",
"\n",
"**Vulnerability:** The `DSA.generate(1024)` function generates a DSA key with a length of 1024 bits, which is considered weak and vulnerable to modern cryptanalysis techniques.\n",
"\n",
"\n",
"**Recommendations:**\n",
"\n",
"*Use a stronger key generation algorithm, such as RSA or ECC, with a key length of at least 2048 bits.\n",
"*Consider using libraries that provide secure defaults for key generation.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"from Crypto.PublicKey import RSA\n",
"\n",
"\n",
"def generate_private_key():\n",
" key = RSA.generate(2048)\n",
" return key.export_key('PEM') # Export in a standard format\n",
"```\n",
"\n",
"### file_name: Bulk/example_101.py\n",
"\n",
"**Vulnerability Name:** Insecure Encryption Mode and Static IV\n",
"\n",
"**Vulnerability:** The code uses AES in CBC mode with a static initialization vector (IV). Using a static IV weakens the encryption and makes it vulnerable to certain attacks, especially when encrypting multiple messages with the same key.\n",
"\n",
"**Recommendations:**\n",
"\n",
"*Use a random IV for each encryption operation.\n",
"*Consider using more secure modes of operation, such as GCM or CTR, which provide authenticated encryption.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"from Crypto.Cipher import AES\n",
"from Crypto.Random import get_random_bytes\n",
"\n",
"def getEncryptionKey(data, key):\n",
" iv\n",
" = get_random_bytes(AES.block_size)\n",
" cipher = AES.new(key, AES.MODE_GCM, iv)\n",
" ciphertext, tag = cipher.encrypt_and_digest(data)\n",
" return iv, ciphertext, tag\n",
"```\n",
"\n",
"### file_name: Bulk/example_102.py\n",
"\n",
"**Vulnerability Name:** Insecure FTP Usage and Hardcoded Credentials\n",
"\n",
"**Vulnerability:** The code uses FTP, which transmits data in plaintext, including usernames and passwords. Additionally, the credentials are hardcoded in the script, making them easily discoverable.\n",
"\n",
"**Recommendations:**\n",
"\n",
"\n",
"*Use a secure file transfer protocol, such as SFTP or FTPS, which encrypts the communication.\n",
"*Avoid hardcoding credentials and instead use environment variables or a secure credential storage mechanism.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"import pysftp # Assuming you have the pysftp\n",
" library installed\n",
"\n",
"def downloadFile(hostname, username, password, sensitive_file):\n",
" with pysftp.Connection(hostname, username=username, password=password) as sftp:\n",
" sftp.get(sensitive_file, preserve_mtime=True)\n",
"```\n",
"\n",
"### file_name:Bulk/example_103.py\n",
"\n",
"**Vulnerability Name:** Potential Issues with Password Hashing\n",
"\n",
"**Vulnerability:** While PBKDF2 is a good choice for password hashing, the code uses a static salt, which reduces its effectiveness against rainbow table attacks.\n",
"\n",
"**Recommendations:**\n",
"\n",
"*Use a unique salt for each password. This can be achieved by generating a random salt for each user and storing it securely alongside the hashed password.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"import os\n",
"\n",
"def getHash(password):\n",
" salt = os.urandom(16) # Generate a random\n",
" salt\n",
" hash = pbkdf2_hmac('sha256', password, salt, 100000)\n",
" return salt, hash # Store the salt with the hash\n",
"```\n",
"\n",
"### file_name: Bulk/example_104.py\n",
"\n",
"**Vulnerability Name:** Hardcoded and Weak Admin Password\n",
"\n",
"**Vulnerability:** The code uses a hardcoded and easily guessable password for the admin user. This makes the system extremely vulnerable to unauthorized access.\n",
"\n",
"**Recommendations:**\n",
"\n",
"*Never hardcode passwords, especially for privileged accounts.\n",
"*Enforce strong password policies that require a minimum length, complexity, and regular password changes.\n",
"*Consider using multi-factor authentication for additional security.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"# Implement a secure password storage and verification mechanism, such as bcrypt or Argon2\n",
"# and avoid hardcoding passwords.\n",
"```\n",
"\n",
"### file_name: Bulk/example_105.py\n",
"\n",
"**Vulnerability Name:** Open Redirect Vulnerability\n",
"\n",
"**Vulnerability:** The `flask_redirect` function redirects the user to any URL provided in the `next` parameter without validation. This can be exploited by attackers to redirect users to malicious websites.\n",
"\n",
"**Recommendations:**\n",
"\n",
"*Validate the `next` parameter to ensure it points to a trusted domain or a specific set of allowed URLs.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"from urllib.parse import urlparse\n",
"\n",
"@app.route('flask_redirect')\n",
"def flask_redirect():\n",
" url = request\n",
".args.get(\"next\")\n",
" parsed_url = urlparse(url)\n",
" if parsed_url.netloc in allowed_domains: # Replace with your allowed domains\n",
" return redirect(url)\n",
" else:\n",
" return \"Invalid redirect URL\"\n",
"```\n",
"\n",
"### file_name: Bulk/example_106.py\n",
"\n",
"**Vulnerability Name:** Potential XXE (XML External Entity) Vulnerability\n",
"\n",
"**Vulnerability:** The code parses XML data from the request without disabling external entity processing. This could allow attackers to inject malicious entities that can lead to information disclosure, denial of service\n",
", or server-side request forgery (SSRF) attacks.\n",
"\n",
"**Recommendations:**\n",
"\n",
"* Disable external entity processing when parsing XML data.\n",
"* Use a library that provides secure XML parsing options.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"from defusedxml import lxml\n",
"\n",
"@app.post(\"/\n",
"upload\")\n",
"def upload():\n",
" xml_src = request.get_data()\n",
" parser = lxml.XMLParser(resolve_entities=False)\n",
" doc = lxml.fromstring(xml_src, parser=parser)\n",
" return lxml.tostring(doc)\n",
"\n",
"```\n",
"\n",
"### file_name: Bulk/example_107.py\n",
"\n",
"**Vulnerability Name:** Potential XXE and Insecure Password Handling\n",
"\n",
"**Vulnerability:** Similar to example_106.py, this code is vulnerable to XXE attacks due to insecure XML parsing. Additionally, the `check\n",
"Password` function is not implemented, and it's unclear how passwords are stored and verified.\n",
"\n",
"**Recommendations:**\n",
"\n",
"*Address the XXE vulnerability as described in example_106.py.\n",
"*Implement a secure password storage and verification mechanism, such as bcrypt or Argon2.\n",
"*Avoid storing passwords in plaintext or using weak hashing algorithms.\n",
"\n",
"**Recommended code:**\n",
"\n",
"```python\n",
"# Implement secure XML parsing as shown in example_106.py\n",
"# Implement a secure password storage and verification mechanism\n",
"```\n",
"\"\"\""
]
},
{
"cell_type": "markdown",
"id": "1tK5AtEdIeep",
"metadata": {
"id": "1tK5AtEdIeep"
},
"source": [
"### Parse the response and export it into a csv for further analysis.\n",
"\n",
"To parse the response we have used regular expression pattern matching. This is easy to configure and maintain\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "yVethA0vhsiv",
"metadata": {
"id": "yVethA0vhsiv"
},
"outputs": [],
"source": [
"import re\n",
"\n",
"import pandas as pd\n",
"\n",
"\n",
"def extract_vulnerability_data(text):\n",
" # Regular expressions for extracting data\n",
" file_pattern = r\"###\\s*file_name:\\s*(.*\\.py)\" # Adjusted for \"file_name\"\n",
" vulnerability_name_pattern = r\"\\*\\*Vulnerability Name:\\*\\*\\s*(.*)\"\n",
" vulnerability_pattern = r\"\\*\\*Vulnerability:\\*\\*\\s*(.*?)(?=\\*\\*Recommendations)\"\n",
" recommendation_pattern = r\"\\*\\*Recommendations:\\*\\*\\s*((?:\\*.*\\n)+)\"\n",
" code_pattern = r\"```python(.*?)```\" # Added pattern for recommended code\n",
"\n",
" data = []\n",
"\n",
" # Iterate through each vulnerability report\n",
" for match in re.finditer(r\"###.*?(?=###|$)\", text, re.DOTALL):\n",
" report = match.group(0)\n",
"\n",
" file_name = re.search(file_pattern, report).group(1)\n",
" vulnerability_name = re.search(vulnerability_name_pattern, report).group(1)\n",
" vulnerability = (\n",
" re.search(vulnerability_pattern, report, re.DOTALL).group(1).strip()\n",
" )\n",
" recommendations = (\n",
" re.search(recommendation_pattern, report, re.DOTALL).group(1).strip()\n",
" )\n",
"\n",
" # Extract recommended code, handling potential absence\n",
" code_match = re.search(code_pattern, report, re.DOTALL)\n",
" recommended_code = code_match.group(1).strip() if code_match else \"N/A\"\n",
"\n",
" data.append(\n",
" {\n",
" \"File Number\": file_name.split(\"_\")[-1].split(\".\")[0],\n",
" \"File Name\": file_name,\n",
" \"Vulnerability Name\": vulnerability_name,\n",
" \"Description\": vulnerability,\n",
" \"Recommendations\": recommendations,\n",
" \"Recommended Code\": recommended_code, # Include recommended code\n",
" }\n",
" )\n",
"\n",
" return pd.DataFrame(data)\n",
"\n",
"\n",
"# Extract data and create DataFrame\n",
"df = extract_vulnerability_data(response_text)\n",
"\n",
"# Display DataFrame in Colab\n",
"display(df)\n",
"\n",
"# Save DataFrame to CSV file\n",
"df.to_csv(\"vulnerability_report_BULK.csv\", index=False)\n",
"\n",
"print(\"CSV file 'vulnerability_report_BULK.csv' created successfully.\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ipY3TXNUieVo",
"metadata": {
"id": "ipY3TXNUieVo"
},
"outputs": [
{
"data": {
"application/javascript": "\n async function download(id, filename, size) {\n if (!google.colab.kernel.accessAllowed) {\n return;\n }\n const div = document.createElement('div');\n const label = document.createElement('label');\n label.textContent = `Downloading \"${filename}\": `;\n div.appendChild(label);\n const progress = document.createElement('progress');\n progress.max = size;\n div.appendChild(progress);\n document.body.appendChild(div);\n\n const buffers = [];\n let downloaded = 0;\n\n const channel = await google.colab.kernel.comms.open(id);\n // Send a message to notify the kernel that we're ready.\n channel.send({})\n\n for await (const message of channel.messages) {\n // Send a message to notify the kernel that we're ready.\n channel.send({})\n if (message.buffers) {\n for (const buffer of message.buffers) {\n buffers.push(buffer);\n downloaded += buffer.byteLength;\n progress.value = downloaded;\n }\n }\n }\n const blob = new Blob(buffers, {type: 'application/binary'});\n const a = document.createElement('a');\n a.href = window.URL.createObjectURL(blob);\n a.download = filename;\n div.appendChild(a);\n a.click();\n div.remove();\n }\n ",
"text/plain": [
"<IPython.core.display.Javascript object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/javascript": "download(\"download_aab82d58-53a1-40c6-be64-e5bb8ff671f5\", \"vulnerability_report_BULK.csv\", 9660)",
"text/plain": [
"<IPython.core.display.Javascript object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"# Display the created CSV file (assuming you have the 'google.colab' library installed)\n",
"from google.colab import files\n",
"\n",
"files.download(\"vulnerability_report_BULK.csv\")"
]
},
{
"cell_type": "markdown",
"id": "7-kVfA_VIVv8",
"metadata": {
"id": "7-kVfA_VIVv8"
},
"source": [
"## Parse the response and export it into JSON output.\n",
"\n",
"NOTE: The following process is just a way to capture the response text and parse it further with regular expression. With Gemini 2.0, we can force the model to respond in JSON structure as well."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "F3hgDQcbn5YO",
"metadata": {
"id": "F3hgDQcbn5YO"
},
"outputs": [],
"source": [
"# Extract data and write to JSON file (assuming you have 'response_text' defined)\n",
"data = extract_vulnerability_data(response_text)\n",
"\n",
"data.to_json(\"vulnerabilities.json\", indent=4)\n",
"\n",
"print(\"Vulnerability data extracted and saved to vulnerabilities.json\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "u-jhc3_Qpx9c",
"metadata": {
"id": "u-jhc3_Qpx9c"
},
"outputs": [
{
"data": {
"application/javascript": "\n async function download(id, filename, size) {\n if (!google.colab.kernel.accessAllowed) {\n return;\n }\n const div = document.createElement('div');\n const label = document.createElement('label');\n label.textContent = `Downloading \"${filename}\": `;\n div.appendChild(label);\n const progress = document.createElement('progress');\n progress.max = size;\n div.appendChild(progress);\n document.body.appendChild(div);\n\n const buffers = [];\n let downloaded = 0;\n\n const channel = await google.colab.kernel.comms.open(id);\n // Send a message to notify the kernel that we're ready.\n channel.send({})\n\n for await (const message of channel.messages) {\n // Send a message to notify the kernel that we're ready.\n channel.send({})\n if (message.buffers) {\n for (const buffer of message.buffers) {\n buffers.push(buffer);\n downloaded += buffer.byteLength;\n progress.value = downloaded;\n }\n }\n }\n const blob = new Blob(buffers, {type: 'application/binary'});\n const a = document.createElement('a');\n a.href = window.URL.createObjectURL(blob);\n a.download = filename;\n div.appendChild(a);\n a.click();\n div.remove();\n }\n ",
"text/plain": [
"<IPython.core.display.Javascript object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/javascript": "download(\"download_7ee96007-491b-4a8f-b52a-6b7d5e63e4b7\", \"vulnerabilities.json\", 11541)",
"text/plain": [
"<IPython.core.display.Javascript object>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"from google.colab import files\n",
"\n",
"files.download(\"vulnerabilities.json\")"
]
},
{
"cell_type": "markdown",
"id": "qbk8xpGJUASo",
"metadata": {
"id": "qbk8xpGJUASo"
},
"source": [
"## Conclusions\n",
"\n",
"In this notebook we have successfully leveraged Gemini 2.0's code scanning and code generation capability to\n",
"1. Analyze multiple python files for potential vulnerabilities\n",
"2. Used Gemini 2.0 to provide recommendations (both in wordings and in the form of code)\n",
"3. Export the details into csv and json for further analysis\n",
"\n",
"The scope of this experiment is limited to identifying issues and providing helpful and contextual modification. Automating remediations or fitting the findings into a review workflow would exist in a more mature tool, and hasn't been considered as part of the experiment. While Gemini 2.0 demonstrates promising capabilities in code analysis, it's important to note that this approach is still experimental. We believe that it is important to explore the potential of this technology for vulnerability detection, and continue development and validation efforts before it can be considered a robust security tool"
]
}
],
"metadata": {
"colab": {
"name": "code_scanning_and_vulnerability_detection.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
gemini/use-cases/code/refactor_legacy_code.ipynb
+1115
View File
File diff suppressed because it is too large Load Diff