src/notebooks/getting_started.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "id": "3027aad445a8ca34",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-04T21:21:47.564156Z",
     "start_time": "2024-07-04T21:21:47.542742Z"
    }
   },
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "cell_type": "code",
   "id": "857d764bfbe99f68",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-04T21:21:47.801729Z",
     "start_time": "2024-07-04T21:21:47.722357Z"
    }
   },
   "source": [
    "import json\n",
    "import logging\n",
    "\n",
    "logging.basicConfig(level=logging.WARNING) # set your preferred logging level here\n",
    "\n",
    "# we import the things we need (obviously)\n",
    "from data.VulnerabilityReport import create_from_flama_json # this is the function that creates the VulnerabilityReport object from the json data\n",
    "from ai.LLM.LLMServiceStrategy import LLMServiceStrategy "
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "id": "f490c5d5b210be2b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-04T21:21:48.170910Z",
     "start_time": "2024-07-04T21:21:47.944920Z"
    }
   },
   "source": [
    "from pathlib import Path\n",
    "\n",
    "data_folder = Path('../../data/') # Please change to the folder, Dataset.json is located in\n",
    "# We need the json data, so we are loading it from the file\n",
    "with open(data_folder / Path('Dataset_Juicebox.json')) as f:\n",
    "    data = json.load(f)"
   ],
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: '../../data/Dataset.json'",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mFileNotFoundError\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[3], line 5\u001B[0m\n\u001B[1;32m      3\u001B[0m data_folder \u001B[38;5;241m=\u001B[39m Path(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124m../../data/\u001B[39m\u001B[38;5;124m'\u001B[39m) \u001B[38;5;66;03m# Please change to the folder, Dataset.json is located in\u001B[39;00m\n\u001B[1;32m      4\u001B[0m \u001B[38;5;66;03m# We need the json data, so we are loading it from the file\u001B[39;00m\n\u001B[0;32m----> 5\u001B[0m \u001B[38;5;28;01mwith\u001B[39;00m \u001B[38;5;28;43mopen\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43mdata_folder\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m/\u001B[39;49m\u001B[43m \u001B[49m\u001B[43mPath\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[38;5;124;43mDataset.json\u001B[39;49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[43m)\u001B[49m\u001B[43m)\u001B[49m \u001B[38;5;28;01mas\u001B[39;00m f:\n\u001B[1;32m      6\u001B[0m     data \u001B[38;5;241m=\u001B[39m json\u001B[38;5;241m.\u001B[39mload(f)\n",
      "File \u001B[0;32m~/Desktop/Uni/SS24/Praktikum/Security-Findings-Recommender-System/src/venv/lib/python3.9/site-packages/IPython/core/interactiveshell.py:310\u001B[0m, in \u001B[0;36m_modified_open\u001B[0;34m(file, *args, **kwargs)\u001B[0m\n\u001B[1;32m    303\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m file \u001B[38;5;129;01min\u001B[39;00m {\u001B[38;5;241m0\u001B[39m, \u001B[38;5;241m1\u001B[39m, \u001B[38;5;241m2\u001B[39m}:\n\u001B[1;32m    304\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mValueError\u001B[39;00m(\n\u001B[1;32m    305\u001B[0m         \u001B[38;5;124mf\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mIPython won\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mt let you open fd=\u001B[39m\u001B[38;5;132;01m{\u001B[39;00mfile\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m by default \u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    306\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mas it is likely to crash IPython. If you know what you are doing, \u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    307\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124myou can use builtins\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124m open.\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    308\u001B[0m     )\n\u001B[0;32m--> 310\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mio_open\u001B[49m\u001B[43m(\u001B[49m\u001B[43mfile\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "\u001B[0;31mFileNotFoundError\u001B[0m: [Errno 2] No such file or directory: '../../data/Dataset.json'"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Choose your Strategy\n",
    "Right now, we support two strategies listed below.\n",
    "Execute the one you like more :)"
   ],
   "id": "133dc5fc298ba3d8"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Option 1) OpenAI - Remote\n",
    "Uses OpenAIs GPT-4o by default.\n",
    "\n",
    "Relatively cheap and quite good.\n"
   ],
   "id": "cbb43a8273ab2e09"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-04T21:21:49.449238Z",
     "start_time": "2024-07-04T21:21:49.228925Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from ai.LLM.Strategies.OpenAIService import OpenAIService\n",
    "\n",
    "my_strategy = OpenAIService()"
   ],
   "id": "a17be348b6b7294a",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Option 2) Anthropic - Remote\n",
    "Uses Anthropics Claude 3 Opus by default.\n",
    "\n",
    "Almost perfect but quite expensive.\n"
   ],
   "id": "33ce7bb1fd20e336"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-04T21:21:50.558588Z",
     "start_time": "2024-07-04T21:21:50.538901Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from ai.LLM.Strategies.OpenAIService import OpenAIService\n",
    "\n",
    "my_strategy = OpenAIService()"
   ],
   "id": "901c5824178df8bb",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Option 3) OLLAMA - Local\n",
    "Here, you can choose the model you want to use. If you don't know which one to choose, you can just leave it empty and the default model will be used. Available models can be found [here](https://ollama.com/library).\n",
    "\n",
    "The performance (both in time and quality) depends heavily on the model used."
   ],
   "id": "3357a17a826c90f7"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-04T21:21:52.159152Z",
     "start_time": "2024-07-04T21:21:51.469910Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from ai.LLM.Strategies.OLLAMAService import OLLAMAService # this is the service that uses ollama to generate solution\n",
    "\n",
    "model_name=None # but can be anything from here https://ollama.com/library\n",
    "my_strategy = OLLAMAService(model_name=model_name)"
   ],
   "id": "845858d36a9b448d",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model llama3:instruct\n"
     ]
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# LLM Service\n",
    "This is the service that uses the strategy you chose. It is the top abstraction layer that will be used to generate solutions."
   ],
   "id": "b70749b293caeb08"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-04T21:21:53.352668Z",
     "start_time": "2024-07-04T21:21:53.332426Z"
    }
   },
   "cell_type": "code",
   "source": "llm_service = LLMServiceStrategy(my_strategy)",
   "id": "9750ed3f60f2a320",
   "outputs": [],
   "execution_count": 7
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Data Setup\n",
    "We instance a vulnerability report object and add findings to it. We also sort the findings by severity.\n",
    "This object takes the LLM service and acts as main abstraction layer for the data.\n"
   ],
   "id": "b7f9f9f8d546dcbb"
  },
  {
   "cell_type": "code",
   "id": "926f7b0460c0f790",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-04T21:21:54.584013Z",
     "start_time": "2024-07-04T21:21:54.531852Z"
    }
   },
   "source": [
    "n=-1 # number of findings to import\n",
    "vulnerability_report = create_from_flama_json(data, n = n, llm_service=llm_service) # here, we create the VulnerabilityReport object, consisting of a list of Finding objects which each have Solution objects. We pass the llm_service to the VulnerabilityReport object, but this can also just be omitted, in which case it will be created in each Finding object.\n",
    "vulnerability_report.sort() # this will sort the findings by severity"
   ],
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'data' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[8], line 2\u001B[0m\n\u001B[1;32m      1\u001B[0m n\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m \u001B[38;5;66;03m# number of findings to import\u001B[39;00m\n\u001B[0;32m----> 2\u001B[0m vulnerability_report \u001B[38;5;241m=\u001B[39m create_from_flama_json(\u001B[43mdata\u001B[49m, n \u001B[38;5;241m=\u001B[39m n, llm_service\u001B[38;5;241m=\u001B[39mllm_service) \u001B[38;5;66;03m# here, we create the VulnerabilityReport object, consisting of a list of Finding objects which each have Solution objects. We pass the llm_service to the VulnerabilityReport object, but this can also just be omitted, in which case it will be created in each Finding object.\u001B[39;00m\n\u001B[1;32m      3\u001B[0m vulnerability_report\u001B[38;5;241m.\u001B[39msort() \u001B[38;5;66;03m# this will sort the findings by severity\u001B[39;00m\n",
      "\u001B[0;31mNameError\u001B[0m: name 'data' is not defined"
     ]
    }
   ],
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "id": "c053459ec572c675",
   "metadata": {},
   "source": "# Lets get going with AI - 🚀🌕"
  },
  {
   "cell_type": "markdown",
   "id": "ecc63b904e7ad48",
   "metadata": {},
   "source": [
    "## First: We add categories to the Findings\n",
    "This will help us later to improve prompts (hopefully)"
   ]
  },
  {
   "cell_type": "code",
   "id": "14edeb4a26eeab5e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-04T21:21:57.189559Z",
     "start_time": "2024-07-04T21:21:57.159896Z"
    }
   },
   "source": [
    "vulnerability_report.add_category() # this will add categories to the findings"
   ],
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'vulnerability_report' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[9], line 1\u001B[0m\n\u001B[0;32m----> 1\u001B[0m \u001B[43mvulnerability_report\u001B[49m\u001B[38;5;241m.\u001B[39madd_category() \u001B[38;5;66;03m# this will add categories to the findings\u001B[39;00m\n",
      "\u001B[0;31mNameError\u001B[0m: name 'vulnerability_report' is not defined"
     ]
    }
   ],
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "id": "bef753bd813064d3",
   "metadata": {},
   "source": [
    "## Then we add solutions 🎉\n",
    "The function below adds long and short solutions to the findings. It also adds keywords for future research.\n",
    "\n",
    "For this function, we have multiple layers that may catch errors. So if you see a warning but no error afterwards, the waring was caught. Only if you see an error, data wasn't generated correctly."
   ]
  },
  {
   "cell_type": "code",
   "id": "d51599fff12405ff",
   "metadata": {},
   "source": [
    "vulnerability_report.add_solution(long=True, short=True, search_term=True) # this will add solutions to the findings"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# Adding unsupervised categories",
   "id": "229f80ff75352254"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "_, Clusterer = vulnerability_report.add_unsupervised_category()\n",
    "Clusterer.get_cluster_graph()"
   ],
   "id": "f813be464bed8a04"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Last but not least: Results\n",
    "We save and display results in the coming cells."
   ],
   "id": "a4ce1d49107d0bab"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "with open(data_folder / Path(f'VulnerabilityReport_{n if n !=-1 else \"all\"}_' + my_strategy.get_model_name().replace(\":\",\"_\") +'.json'), 'w') as f:\n",
    "    f.write(str(json.dumps(vulnerability_report.to_dict())))"
   ],
   "id": "73d0454da9633cc0",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "c827dd06",
   "metadata": {},
   "source": [
    "from IPython.display import HTML, Markdown\n",
    "Markdown(vulnerability_report.to_html(table=True))\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "HTML(str(vulnerability_report).replace('\\n', '<br />'))\n",
   "id": "496eb3b60477c4b4",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "c1a753a332bd90d7",
   "outputs": [],
   "execution_count": null
  }
 ],
 "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}