From a023d406505a8e33252cd9846c07209ff7f2f98e Mon Sep 17 00:00:00 2001 From: ErdemOzgen Date: Thu, 30 Nov 2023 11:25:37 +0300 Subject: [PATCH 01/18] Add ollama localgpt to report generator and AttackSuggestionGenerator classes --- web/reNgine/gpt.py | 151 +++++++++++++++++++++++---------------------- 1 file changed, 76 insertions(+), 75 deletions(-) diff --git a/web/reNgine/gpt.py b/web/reNgine/gpt.py index 4bae2563..40c7fe2a 100644 --- a/web/reNgine/gpt.py +++ b/web/reNgine/gpt.py @@ -2,13 +2,16 @@ import openai import re from reNgine.common_func import get_open_ai_key, extract_between from reNgine.definitions import VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE, ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT +from langchain.llms import Ollama class GPTVulnerabilityReportGenerator: def __init__(self): self.api_key = get_open_ai_key() self.model_name = 'gpt-3.5-turbo' - + if not self.api_key: + self.ollama = Ollama(base_url='http://localhost:11434', model="llama2-uncensored") + def get_vulnerability_description(self, description): """Generate Vulnerability Description using GPT. @@ -24,93 +27,91 @@ class GPTVulnerabilityReportGenerator: } """ if not self.api_key: - return { - 'status': False, - 'error': 'No OpenAI keys provided.' - } - openai.api_key = self.api_key - try: - gpt_response = openai.ChatCompletion.create( - model=self.model_name, - messages=[ - {'role': 'system', 'content': VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE}, - {'role': 'user', 'content': description} - ] - ) + prompt = ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT + "\nUser: " + input + response_content = self.ollama(prompt) + else: + openai.api_key = self.api_key + try: + gpt_response = openai.ChatCompletion.create( + model=self.model_name, + messages=[ + {'role': 'system', 'content': VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE}, + {'role': 'user', 'content': description} + ] + ) - response_content = gpt_response['choices'][0]['message']['content'] + response_content = gpt_response['choices'][0]['message']['content'] + except Exception as e: + return { + 'status': False, + 'error': str(e) + } + vuln_description_pattern = re.compile( + r"[Vv]ulnerability [Dd]escription:(.*?)(?:\n\n[Ii]mpact:|$)", + re.DOTALL + ) + impact_pattern = re.compile( + r"[Ii]mpact:(.*?)(?:\n\n[Rr]emediation:|$)", + re.DOTALL + ) + remediation_pattern = re.compile( + r"[Rr]emediation:(.*?)(?:\n\n[Rr]eferences:|$)", + re.DOTALL + ) - vuln_description_pattern = re.compile( - r"[Vv]ulnerability [Dd]escription:(.*?)(?:\n\n[Ii]mpact:|$)", - re.DOTALL - ) - impact_pattern = re.compile( - r"[Ii]mpact:(.*?)(?:\n\n[Rr]emediation:|$)", - re.DOTALL - ) - remediation_pattern = re.compile( - r"[Rr]emediation:(.*?)(?:\n\n[Rr]eferences:|$)", - re.DOTALL - ) + description_section = extract_between(response_content, vuln_description_pattern) + impact_section = extract_between(response_content, impact_pattern) + remediation_section = extract_between(response_content, remediation_pattern) + references_start_index = response_content.find("References:") + references_section = response_content[references_start_index + len("References:"):].strip() - description_section = extract_between(response_content, vuln_description_pattern) - impact_section = extract_between(response_content, impact_pattern) - remediation_section = extract_between(response_content, remediation_pattern) - references_start_index = response_content.find("References:") - references_section = response_content[references_start_index + len("References:"):].strip() - - url_pattern = re.compile(r'https://\S+') - urls = url_pattern.findall(references_section) - - return { - 'status': True, - 'description': description_section, - 'impact': impact_section, - 'remediation': remediation_section, - 'references': urls, - } - except Exception as e: - return { - 'status': False, - 'error': str(e) - } + url_pattern = re.compile(r'https://\S+') + urls = url_pattern.findall(references_section) + return { + 'status': True, + 'description': description_section, + 'impact': impact_section, + 'remediation': remediation_section, + 'references': urls, + } class GPTAttackSuggestionGenerator: def __init__(self): self.api_key = get_open_ai_key() self.model_name = 'gpt-3.5-turbo' + if not self.api_key: + self.ollama = Ollama(base_url='http://localhost:11434', model="llama2-uncensored") def get_attack_suggestion(self, input): ''' input (str): input for gpt ''' if not self.api_key: - return { - 'status': False, - 'error': 'No OpenAI keys provided.', - 'input': input - } - openai.api_key = self.api_key - print(input) - try: - gpt_response = openai.ChatCompletion.create( - model=self.model_name, - messages=[ - {'role': 'system', 'content': ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT}, - {'role': 'user', 'content': input} - ] - ) - response_content = gpt_response['choices'][0]['message']['content'] - return { - 'status': True, - 'description': response_content, - 'input': input - } - except Exception as e: - return { - 'status': False, - 'error': str(e), - 'input': input - } + prompt = ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT + "\nUser: " + input + response_content = self.ollama(prompt) + else: + openai.api_key = self.api_key + print(input) + try: + gpt_response = openai.ChatCompletion.create( + model=self.model_name, + messages=[ + {'role': 'system', 'content': ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT}, + {'role': 'user', 'content': input} + ] + ) + response_content = gpt_response['choices'][0]['message']['content'] + except Exception as e: + return { + 'status': False, + 'error': str(e), + 'input': input + } + return { + 'status': True, + 'description': response_content, + 'input': input + } + \ No newline at end of file From 923873550634916be1277f8c1366c1c0f0067dd5 Mon Sep 17 00:00:00 2001 From: ErdemOzgen Date: Thu, 30 Nov 2023 11:26:30 +0300 Subject: [PATCH 02/18] Add Erdem Ozgen to CONTRIBUTORS.md --- CONTRIBUTORS.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 00cf16f0..297b9193 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -8,6 +8,7 @@ Thanks to these individuals for making reNgine awesome by fixing bugs, resolving * [Suprita-25](https://github.com/Suprita-25) * [TheBinitGhimire](https://github.com/TheBinitGhimire) * [Vinay Leo](https://github.com/vinaynm) +* [Erdem Ozgen](https://github.com/ErdemOzgen) *If you have created a Pull request, feel free to add your name here, because we know you are awesome and deserve thanks from the community!* From 6cbbe0bbe3c90de27c3d9379f46597c7c681ee0d Mon Sep 17 00:00:00 2001 From: ErdemOzgen Date: Thu, 30 Nov 2023 13:16:44 +0300 Subject: [PATCH 03/18] Add langchain library --- web/requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/web/requirements.txt b/web/requirements.txt index 08c2d49c..e402d590 100644 --- a/web/requirements.txt +++ b/web/requirements.txt @@ -37,3 +37,4 @@ whatportis weasyprint==53.3 wafw00f==2.2.0 xmltodict==0.13.0 +langchain==0.0.343 From 266f8983935d10f841f5f0bb7d30c0d407d7d369 Mon Sep 17 00:00:00 2001 From: ErdemOzgen Date: Thu, 30 Nov 2023 15:59:00 +0300 Subject: [PATCH 04/18] Add ollama dockerservice --- docker-compose.yml | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/docker-compose.yml b/docker-compose.yml index 3d2dc052..0618f534 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -143,7 +143,17 @@ services: - scan_results:/usr/src/scan_results networks: - rengine_network - + ollama: + image: ollama/ollama + container_name: ollama + volumes: + - ollama_data:/root/.ollama + ports: + - "11434:11434" + networks: + - rengine_network + restart: always + command: ["ollama", "run", "llama2-uncensored"] networks: rengine_network: @@ -157,6 +167,7 @@ volumes: wordlist: scan_results: static_volume: + ollama_data: secrets: proxy.ca: From 75ebc3db48f52dff198740f7b470ad5d9735133c Mon Sep 17 00:00:00 2001 From: ErdemOzgen Date: Thu, 30 Nov 2023 16:06:21 +0300 Subject: [PATCH 05/18] Add ollama service to docker-compose --- Makefile | 2 +- docker-compose.dev.yml | 12 ++++++++++++ web/reNgine/gpt.py | 4 ++-- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index f1559dcf..8e23217d 100644 --- a/Makefile +++ b/Makefile @@ -7,7 +7,7 @@ include .env COMPOSE_PREFIX_CMD := COMPOSE_DOCKER_CLI_BUILD=1 COMPOSE_ALL_FILES := -f docker-compose.yml -SERVICES := db web proxy redis celery celery-beat +SERVICES := db web proxy redis celery celery-beat ollama # -------------------------- diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index ffd733e3..78b32418 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -115,6 +115,17 @@ services: - celery-beat networks: - rengine_network + ollama: + image: ollama/ollama + container_name: ollama + volumes: + - ollama_data:/root/.ollama + ports: + - "11434:11434" + networks: + - rengine_network + restart: always + command: ["ollama", "run", "llama2-uncensored"] networks: rengine_network: @@ -126,3 +137,4 @@ volumes: github_repos: wordlist: scan_results: + ollama_data: diff --git a/web/reNgine/gpt.py b/web/reNgine/gpt.py index 40c7fe2a..2e6682ff 100644 --- a/web/reNgine/gpt.py +++ b/web/reNgine/gpt.py @@ -10,7 +10,7 @@ class GPTVulnerabilityReportGenerator: self.api_key = get_open_ai_key() self.model_name = 'gpt-3.5-turbo' if not self.api_key: - self.ollama = Ollama(base_url='http://localhost:11434', model="llama2-uncensored") + self.ollama = Ollama(base_url='http://ollama:11434', model="llama2-uncensored") def get_vulnerability_description(self, description): """Generate Vulnerability Description using GPT. @@ -82,7 +82,7 @@ class GPTAttackSuggestionGenerator: self.api_key = get_open_ai_key() self.model_name = 'gpt-3.5-turbo' if not self.api_key: - self.ollama = Ollama(base_url='http://localhost:11434', model="llama2-uncensored") + self.ollama = Ollama(base_url='http://ollama:11434', model="llama2-uncensored") def get_attack_suggestion(self, input): ''' From bc4f40784b9a6804aa1c0636aaa0b80bba1c4c4e Mon Sep 17 00:00:00 2001 From: ErdemOzgen Date: Thu, 30 Nov 2023 21:55:14 +0300 Subject: [PATCH 06/18] Add Docker command to run llama2-uncensored --- docker-compose.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index 0618f534..b22c6c67 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -152,8 +152,6 @@ services: - "11434:11434" networks: - rengine_network - restart: always - command: ["ollama", "run", "llama2-uncensored"] networks: rengine_network: From d8000812051cc0afd5ffe9d20acbdcab9d6ed3e5 Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sat, 20 Apr 2024 16:18:21 +0530 Subject: [PATCH 07/18] remove exposed port and command --- docker-compose.dev.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 78b32418..2afaf3dd 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -120,12 +120,12 @@ services: container_name: ollama volumes: - ollama_data:/root/.ollama - ports: - - "11434:11434" + # ports: + # - "11434:11434" networks: - rengine_network restart: always - command: ["ollama", "run", "llama2-uncensored"] + # command: ["ollama", "run", "llama2-uncensored"] networks: rengine_network: From 36967f8fc3d9e0da26de3d2f034abafd5cc5cd23 Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sat, 20 Apr 2024 16:18:59 +0530 Subject: [PATCH 08/18] added custom client for ollama --- web/reNgine/definitions.py | 6 ++++++ web/reNgine/llm.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) create mode 100644 web/reNgine/llm.py diff --git a/web/reNgine/definitions.py b/web/reNgine/definitions.py index 7034fbca..5abb5ebf 100644 --- a/web/reNgine/definitions.py +++ b/web/reNgine/definitions.py @@ -436,6 +436,12 @@ PERM_INITATE_SCANS_SUBSCANS = 'initiate_scans_subscans' FOUR_OH_FOUR_URL = '/404/' +############################################################################### +# OLLAMA DEFINITIONS +############################################################################### +OLLAMA_INSTANCE = 'http://ollama:11434' + + # GPT Vulnerability Report Generator VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE = """ You are a highly skilled penetration tester who has recently completed a penetration testing. diff --git a/web/reNgine/llm.py b/web/reNgine/llm.py new file mode 100644 index 00000000..26adb9ff --- /dev/null +++ b/web/reNgine/llm.py @@ -0,0 +1,28 @@ +from reNgine.definitions import OLLAMA_INSTANCE +from ollama import Client + +class CustomOllamaClient: + + def __init__(self, ): + self.instance_url = OLLAMA_INSTANCE + self.client = None + + def connect(self): + try: + self.client = Client( + host=self.instance_url + ) + return { + 'status': True + } + except Exception as e: + return { + 'status': False, + 'error': str(e) + } + + def list_models(self,): + return self.client.list() + + def pull_model(self, model_name): + return self.client.pull(model_name) \ No newline at end of file From 034c425b1b9b5f941055d6a45222fd333eea1ce3 Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sat, 20 Apr 2024 16:20:51 +0530 Subject: [PATCH 09/18] add ui for llm toolkit --- web/reNgine/llm.py | 1 + .../templates/scanEngine/settings/api.html | 8 +- .../scanEngine/settings/llm_toolkit.html | 76 +++++++++++++++++++ web/scanEngine/urls.py | 4 + web/scanEngine/views.py | 12 +++ web/templates/base/_items/top_nav.html | 1 + 6 files changed, 98 insertions(+), 4 deletions(-) create mode 100644 web/scanEngine/templates/scanEngine/settings/llm_toolkit.html diff --git a/web/reNgine/llm.py b/web/reNgine/llm.py index 26adb9ff..43b3eda6 100644 --- a/web/reNgine/llm.py +++ b/web/reNgine/llm.py @@ -22,6 +22,7 @@ class CustomOllamaClient: } def list_models(self,): + models = self.client.list() return self.client.list() def pull_model(self, model_name): diff --git a/web/scanEngine/templates/scanEngine/settings/api.html b/web/scanEngine/templates/scanEngine/settings/api.html index f4a2d6a4..150e912a 100644 --- a/web/scanEngine/templates/scanEngine/settings/api.html +++ b/web/scanEngine/templates/scanEngine/settings/api.html @@ -2,7 +2,7 @@ {% load static %} {% load custom_tags %} {% block title %} -reNgine Settings +API Vault {% endblock title %} {% block custom_js_css_link %} @@ -27,8 +27,8 @@ API Vault
- -

OpenAI keys will be used to generate vulnerability description, remediation, impact and vulnerability report writing using ChatGPT.

+ +

OpenAI keys will be used to generate vulnerability description, remediation, impact and vulnerability report writing using GPT.

{% if openai_key %} {% else %} @@ -38,7 +38,7 @@ API Vault
-

Netlas keys will be used to get whois information and other OSINT data.

+

Netlas keys will be used to get whois information and other OSINT related data.

{% if netlas_key %} {% else %} diff --git a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html new file mode 100644 index 00000000..5d60c36e --- /dev/null +++ b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html @@ -0,0 +1,76 @@ +{% extends 'base/base.html' %} +{% load static %} +{% load custom_tags %} +{% block title %} +LLM Toolkit +{% endblock title %} + +{% block custom_js_css_link %} +{% endblock custom_js_css_link %} + +{% block breadcrumb_title %} +
  • Settings
    • +
  • LLM Toolkit
    • +{% endblock breadcrumb_title %} + +{% block page_title %} +LLM Toolkit +{% endblock page_title %} + +{% block main_content %} +
    +
    +
    +
    +

    reNgine makes use of various LLMs to enhance reporting process. Using various LLM AI Models penetration testers will be able to to generate detailed, insightful penetration testing reports. +
    + LLM Toolkit includes the ability to download new LLMs, view available models, and delete models no longer needed, and also choose between various models. +

    +
    +
    +
    +
    {{installed_models|len}} Installed LLM Models
    +
    + {% for model in installed_models %} +
    +
    +
    +
    + +
    + Delete +
    +
    +

    + {{model.name}} +

    +

    + +

    + + + {{model.modified_at}} Modified at + +
    + + + {{model.details.parameter_size}} Parameters + + + + {{model.details.family}} Family + +

    +
    +
    +
    + {% endfor %} +
    +
    +{% endblock main_content %} + + +{% block page_level_script %} +{% endblock page_level_script %} diff --git a/web/scanEngine/urls.py b/web/scanEngine/urls.py index ff7d2812..89c5084b 100644 --- a/web/scanEngine/urls.py +++ b/web/scanEngine/urls.py @@ -56,6 +56,10 @@ urlpatterns = [ '/tool_arsenal', views.tool_arsenal_section, name='tool_arsenal'), + path( + '/llm_toolkit', + views.llm_toolkit_section, + name='llm_toolkit'), path( '/rengine_settings', views.rengine_settings, diff --git a/web/scanEngine/views.py b/web/scanEngine/views.py index 561006c0..42eb2fc8 100644 --- a/web/scanEngine/views.py +++ b/web/scanEngine/views.py @@ -11,6 +11,7 @@ from django.urls import reverse from rolepermissions.decorators import has_permission_decorator from reNgine.common_func import * +from reNgine.llm import CustomOllamaClient from reNgine.tasks import (run_command, send_discord_message, send_slack_message, send_telegram_message) from scanEngine.forms import * from scanEngine.forms import ConfigurationForm @@ -456,6 +457,17 @@ def tool_arsenal_section(request, slug): return render(request, 'scanEngine/settings/tool_arsenal.html', context) +@has_permission_decorator(PERM_MODIFY_SYSTEM_CONFIGURATIONS, redirect_url=FOUR_OH_FOUR_URL) +def llm_toolkit_section(request, slug): + context = {} + ollama = CustomOllamaClient() + res = ollama.connect() + if not res.get('status'): + + context['installed_models'] = ollama.list_models() + return render(request, 'scanEngine/settings/llm_toolkit.html', context) + + @has_permission_decorator(PERM_MODIFY_SYSTEM_CONFIGURATIONS, redirect_url=FOUR_OH_FOUR_URL) def api_vault(request, slug): context = {} diff --git a/web/templates/base/_items/top_nav.html b/web/templates/base/_items/top_nav.html index dff468eb..9f9c9275 100644 --- a/web/templates/base/_items/top_nav.html +++ b/web/templates/base/_items/top_nav.html @@ -81,6 +81,7 @@ {% endif %} {% if user|can:'modify_system_configurations' %} API Vault + LLM Toolkit Tools Arsenal {% endif %} {% if user|can:'modify_scan_report' %} From 2e48082b46d9a307964aea89efd871e5d8b7114a Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sat, 20 Apr 2024 17:34:21 +0530 Subject: [PATCH 10/18] add new llm button --- web/reNgine/llm.py | 3 +-- .../templates/scanEngine/settings/llm_toolkit.html | 12 ++++++++++-- web/scanEngine/views.py | 2 +- 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/web/reNgine/llm.py b/web/reNgine/llm.py index 43b3eda6..efa2f408 100644 --- a/web/reNgine/llm.py +++ b/web/reNgine/llm.py @@ -22,8 +22,7 @@ class CustomOllamaClient: } def list_models(self,): - models = self.client.list() - return self.client.list() + return self.client.list().get('models', []) def pull_model(self, model_name): return self.client.pull(model_name) \ No newline at end of file diff --git a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html index 5d60c36e..638afdab 100644 --- a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html +++ b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html @@ -29,7 +29,12 @@ LLM Toolkit
    -
    {{installed_models|len}} Installed LLM Models
    +
    +
    + Add New LLM +
    +
    +
    {{installed_models|length}} Installed LLM Models
    {% for model in installed_models %}
    @@ -47,7 +52,6 @@ LLM Toolkit {{model.name}}

    -

    @@ -66,6 +70,10 @@ LLM Toolkit

    + {% empty %} + {% endfor %} diff --git a/web/scanEngine/views.py b/web/scanEngine/views.py index 42eb2fc8..ec2f2b57 100644 --- a/web/scanEngine/views.py +++ b/web/scanEngine/views.py @@ -463,7 +463,7 @@ def llm_toolkit_section(request, slug): ollama = CustomOllamaClient() res = ollama.connect() if not res.get('status'): - + pass context['installed_models'] = ollama.list_models() return render(request, 'scanEngine/settings/llm_toolkit.html', context) From 78bb203492a0ec58a6cdd5783fdc13a1329e8726 Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sun, 21 Apr 2024 08:30:02 +0530 Subject: [PATCH 11/18] added ollama manager to pull and delete models --- web/api/urls.py | 4 ++++ web/api/views.py | 53 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/web/api/urls.py b/web/api/urls.py index f3a09b54..f4ee867e 100644 --- a/web/api/urls.py +++ b/web/api/urls.py @@ -178,6 +178,10 @@ urlpatterns = [ 'tool/uninstall/', UninstallTool.as_view(), name='uninstall_tool'), + path( + 'tool/ollama/', + OllamaManager.as_view(), + name='ollama_manager'), path( 'rengine/update/', RengineUpdateCheck.as_view(), diff --git a/web/api/views.py b/web/api/views.py index c81bd149..5f1f0def 100644 --- a/web/api/views.py +++ b/web/api/views.py @@ -33,6 +33,59 @@ from .serializers import * logger = logging.getLogger(__name__) +class OllamaManager(APIView): + def get(self, request): + """ + API to download Ollama Models + sends a POST request to download the model + """ + req = self.request + model_name = req.query_params.get('model') + response = { + 'status': False + } + try: + pull_model_api = f'{OLLAMA_INSTANCE}/api/pull' + _response = requests.post( + pull_model_api, + json={ + 'name': model_name, + 'stream': False + } + ).json() + if _response.get('error'): + response['status'] = False + response['error'] = _response.get('error') + else: + response['status'] = True + except Exception as e: + response['error'] = str(e) + return Response(response) + + def delete(self, request): + req = self.request + model_name = req.query_params.get('model') + delete_model_api = f'{OLLAMA_INSTANCE}/api/delete' + response = { + 'status': False + } + try: + _response = requests.delete( + delete_model_api, + json={ + 'name': model_name + } + ).json() + if _response.get('error'): + response['status'] = False + response['error'] = _response.get('error') + else: + response['status'] = True + except Exception as e: + response['error'] = str(e) + return Response(response) + + class GPTAttackSuggestion(APIView): def get(self, request): req = self.request From 474bb25ae0ba017244895b2e17d6e948693a09c4 Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sun, 21 Apr 2024 08:30:48 +0530 Subject: [PATCH 12/18] display available models in ui and create delete model js function --- web/reNgine/llm.py | 28 ------------ .../scanEngine/settings/llm_toolkit.html | 43 +++++++++++++++++-- web/scanEngine/views.py | 18 +++++--- 3 files changed, 51 insertions(+), 38 deletions(-) delete mode 100644 web/reNgine/llm.py diff --git a/web/reNgine/llm.py b/web/reNgine/llm.py deleted file mode 100644 index efa2f408..00000000 --- a/web/reNgine/llm.py +++ /dev/null @@ -1,28 +0,0 @@ -from reNgine.definitions import OLLAMA_INSTANCE -from ollama import Client - -class CustomOllamaClient: - - def __init__(self, ): - self.instance_url = OLLAMA_INSTANCE - self.client = None - - def connect(self): - try: - self.client = Client( - host=self.instance_url - ) - return { - 'status': True - } - except Exception as e: - return { - 'status': False, - 'error': str(e) - } - - def list_models(self,): - return self.client.list().get('models', []) - - def pull_model(self, model_name): - return self.client.pull(model_name) \ No newline at end of file diff --git a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html index 638afdab..ef82a728 100644 --- a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html +++ b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html @@ -1,6 +1,6 @@ {% extends 'base/base.html' %} {% load static %} -{% load custom_tags %} +{% load humanize %} {% block title %} LLM Toolkit {% endblock title %} @@ -31,7 +31,7 @@ LLM Toolkit
    - Add New LLM + Add new tool
    {{installed_models|length}} Installed LLM Models
    @@ -45,7 +45,7 @@ LLM Toolkit
    - Delete + Delete

    @@ -55,7 +55,7 @@ LLM Toolkit

    - {{model.modified_at}} Modified at + Modified {{model.modified_at|naturaltime}}
    @@ -67,6 +67,7 @@ LLM Toolkit {{model.details.family}} Family

    + @@ -81,4 +82,38 @@ LLM Toolkit {% block page_level_script %} + {% endblock page_level_script %} diff --git a/web/scanEngine/views.py b/web/scanEngine/views.py index ec2f2b57..e1932442 100644 --- a/web/scanEngine/views.py +++ b/web/scanEngine/views.py @@ -4,6 +4,7 @@ import re import shutil import subprocess +from datetime import datetime from django import http from django.contrib import messages from django.shortcuts import get_object_or_404, render @@ -11,7 +12,6 @@ from django.urls import reverse from rolepermissions.decorators import has_permission_decorator from reNgine.common_func import * -from reNgine.llm import CustomOllamaClient from reNgine.tasks import (run_command, send_discord_message, send_slack_message, send_telegram_message) from scanEngine.forms import * from scanEngine.forms import ConfigurationForm @@ -460,11 +460,17 @@ def tool_arsenal_section(request, slug): @has_permission_decorator(PERM_MODIFY_SYSTEM_CONFIGURATIONS, redirect_url=FOUR_OH_FOUR_URL) def llm_toolkit_section(request, slug): context = {} - ollama = CustomOllamaClient() - res = ollama.connect() - if not res.get('status'): - pass - context['installed_models'] = ollama.list_models() + list_all_models_url = f'{OLLAMA_INSTANCE}/api/tags' + response = requests.get(list_all_models_url) + if response.status_code == 200: + models = response.json() + models = models.get('models') + date_format = "%Y-%m-%dT%H:%M:%S" + models = [{**model, 'modified_at': datetime.strptime(model['modified_at'].split('.')[0], date_format)} for model in models] + context['installed_models'] = models + print(context['installed_models']) + else: + context['installed_models'] = [] return render(request, 'scanEngine/settings/llm_toolkit.html', context) From b8e99205b3ef33aaf937d03ddb13f168c8bbd402 Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sun, 21 Apr 2024 08:57:55 +0530 Subject: [PATCH 13/18] added add model and delete model ui and endpoints --- .../scanEngine/settings/llm_toolkit.html | 101 +++++++++++++++++- 1 file changed, 99 insertions(+), 2 deletions(-) diff --git a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html index ef82a728..702ee047 100644 --- a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html +++ b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html @@ -31,7 +31,7 @@ LLM Toolkit
    - Add new tool + Add new model
    {{installed_models|length}} Installed LLM Models
    @@ -105,7 +105,26 @@ LLM Toolkit }).then(function(response) { return response.json(); }).then(function(data) { - return location.reload(); + if (data.status){ + swal.insertQueueStep({ + icon: 'error', + title: 'Oops! Unable to delete the model!' + }) + } + else{ + swal.queue([{ + title: 'Model Successfully deleted!', + icon: 'success', + showCancelButton: false, + confirmButtonText: 'Okay', + padding: '2em', + showLoaderOnConfirm: true, + preConfirm: function() { + location.reload(); + } + }]); + } + //return location.reload(); }).catch(function() { swal.insertQueueStep({ icon: 'error', @@ -115,5 +134,83 @@ LLM Toolkit } }]) } + + function showAddNewModelModal(){ + $('#modal_title').html('Add new LLM Model'); + $('#modal-content').empty(); + $('#modal-content').append(` +

    You can find the list of supported models in Ollama Library

    +

    We recommend using llama2-uncensored model for better results.

    +
    + + +
    +
    + +
    + `); + $('#modal_dialog').modal('show'); + } + + function download_model(){ + var model_name = $('#model_name').val(); + if (model_name == ""){ + Swal.fire({ + title: 'Oops!', + text: 'Model name is required', + icon: 'error' + }); + return; + } + var url = "/api/tool/ollama/?model=" + model_name; + swal.queue([{ + title: 'Are you sure you want to download this model?', + text: "Downloading models can take a long time, sometimes a few minutes. Please be patient.", + icon: 'info', + showCancelButton: true, + confirmButtonText: 'Download', + padding: '2em', + showLoaderOnConfirm: true, + preConfirm: function() { + return fetch(url, { + method: 'GET', + credentials: "same-origin", + headers: { + "X-CSRFToken": getCookie("csrftoken") + } + }).then(function(response) { + return response.json(); + }).then(function(data) { + if (!data.status){ + swal.insertQueueStep({ + icon: 'error', + title: 'Oops! Unable to download the model, model does not exist!' + }) + } + else{ + swal.queue([{ + title: 'Model Successfully downloaded!', + icon: 'success', + showCancelButton: false, + confirmButtonText: 'Okay', + padding: '2em', + showLoaderOnConfirm: true, + preConfirm: function() { + location.reload(); + } + }]); + } + //return location.reload(); + }).catch(function() { + swal.insertQueueStep({ + icon: 'error', + title: 'Oops! Unable to download the model!' + }) + }) + } + }]) + } + + {% endblock page_level_script %} From 1f35de39203c929ba1f529c5091f9b6a1add1af1 Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sun, 21 Apr 2024 10:31:00 +0530 Subject: [PATCH 14/18] added put method to select model --- web/api/views.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/web/api/views.py b/web/api/views.py index 5f1f0def..fe6cd6cf 100644 --- a/web/api/views.py +++ b/web/api/views.py @@ -84,6 +84,22 @@ class OllamaManager(APIView): except Exception as e: response['error'] = str(e) return Response(response) + + def put(self, request): + req = self.request + model_name = req.query_params.get('model') + response = { + 'status': False + } + try: + OllamaSettings.objects.update_or_create( + defaults={'selected_model': model_name}, + id=1 + ) + response['status'] = True + except Exception as e: + response['error'] = str(e) + return Response(response) class GPTAttackSuggestion(APIView): From 5ab92fcda9f7f3779122c59012a8316a9bd904a7 Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sun, 21 Apr 2024 10:31:31 +0530 Subject: [PATCH 15/18] Added ollama settings to store selected models --- .../migrations/0010_ollamasettings.py | 20 +++++ web/dashboard/models.py | 8 ++ web/reNgine/definitions.py | 40 ++++++++++ .../scanEngine/settings/llm_toolkit.html | 73 +++++++++++++++++-- web/scanEngine/views.py | 23 ++++-- 5 files changed, 153 insertions(+), 11 deletions(-) create mode 100644 web/dashboard/migrations/0010_ollamasettings.py diff --git a/web/dashboard/migrations/0010_ollamasettings.py b/web/dashboard/migrations/0010_ollamasettings.py new file mode 100644 index 00000000..1e47e5d0 --- /dev/null +++ b/web/dashboard/migrations/0010_ollamasettings.py @@ -0,0 +1,20 @@ +# Generated by Django 3.2.4 on 2024-04-21 04:35 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('dashboard', '0009_delete_openaikeys'), + ] + + operations = [ + migrations.CreateModel( + name='OllamaSettings', + fields=[ + ('id', models.AutoField(primary_key=True, serialize=False)), + ('selected_model', models.CharField(max_length=500)), + ], + ), + ] diff --git a/web/dashboard/models.py b/web/dashboard/models.py index 3d600174..58bb2183 100644 --- a/web/dashboard/models.py +++ b/web/dashboard/models.py @@ -24,6 +24,14 @@ class OpenAiAPIKey(models.Model): def __str__(self): return self.key + + +class OllamaSettings(models.Model): + id = models.AutoField(primary_key=True) + selected_model = models.CharField(max_length=500) + + def __str__(self): + return self.selected_model class NetlasAPIKey(models.Model): diff --git a/web/reNgine/definitions.py b/web/reNgine/definitions.py index 5abb5ebf..d4abe300 100644 --- a/web/reNgine/definitions.py +++ b/web/reNgine/definitions.py @@ -441,6 +441,46 @@ FOUR_OH_FOUR_URL = '/404/' ############################################################################### OLLAMA_INSTANCE = 'http://ollama:11434' +DEFAULT_GPT_MODELS = [ + { + 'name': 'gpt-3', + 'model': 'gpt-3', + 'modified_at': '', + 'details': { + 'family': 'GPT', + 'parameter_size': '~175B', + } + }, + { + 'name': 'gpt-3.5-turbo', + 'model': 'gpt-3.5-turbo', + 'modified_at': '', + 'details': { + 'family': 'GPT', + 'parameter_size': '~7B', + } + }, + { + 'name': 'gpt-4', + 'model': 'gpt-4', + 'modified_at': '', + 'details': { + 'family': 'GPT', + 'parameter_size': '~1.7T', + } + }, + { + 'name': 'gpt-4-turbo', + 'model': 'gpt-4', + 'modified_at': '', + 'details': { + 'family': 'GPT', + 'parameter_size': '~1.7T', + } + } +] + + # GPT Vulnerability Report Generator VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE = """ diff --git a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html index 702ee047..b6c87c7e 100644 --- a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html +++ b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html @@ -34,7 +34,7 @@ LLM Toolkit Add new model -
    {{installed_models|length}} Installed LLM Models
    +
    {{installed_models|length}} available Models
    {% for model in installed_models %}
    @@ -45,17 +45,31 @@ LLM Toolkit
    - Delete + {% if model.is_local %} + Delete + {% endif %} + {% if not model.selected %} + Use Model + {% endif %}

    - {{model.name}} + {{model.name}} {% if model.selected %}Selected Model{% endif %}

    - Modified {{model.modified_at|naturaltime}} + Modified {% if model.modified_at %}{{model.modified_at|naturaltime}} {% else %} NA{% endif %} + +
    + + + {% if model.is_local %} + Locally installed model + {% else %} + Open AI Model + {% endif %}
    @@ -67,7 +81,6 @@ LLM Toolkit {{model.details.family}} Family

    -
    @@ -210,7 +223,55 @@ LLM Toolkit } }]) } - + function selectModel(model_name){ + var url = "/api/tool/ollama/?model=" + model_name; + swal.queue([{ + title: 'Are you sure you want to select this model?', + text: "This model will be used to generate Scan Reports and Attack Suggestions.", + icon: 'info', + showCancelButton: true, + confirmButtonText: 'Select', + padding: '2em', + showLoaderOnConfirm: true, + preConfirm: function() { + return fetch(url, { + method: 'PUT', + credentials: "same-origin", + headers: { + "X-CSRFToken": getCookie("csrftoken") + } + }).then(function(response) { + return response.json(); + }).then(function(data) { + if (!data.status){ + swal.insertQueueStep({ + icon: 'error', + title: 'Oops! Unable to select the model!' + }) + } + else{ + swal.queue([{ + title: 'Model Successfully selected!', + icon: 'success', + showCancelButton: false, + confirmButtonText: 'Okay', + padding: '2em', + showLoaderOnConfirm: true, + preConfirm: function() { + location.reload(); + } + }]); + } + //return location.reload(); + }).catch(function() { + swal.insertQueueStep({ + icon: 'error', + title: 'Oops! Unable to select the model!' + }) + }) + } + }]) + } {% endblock page_level_script %} diff --git a/web/scanEngine/views.py b/web/scanEngine/views.py index e1932442..b4f19574 100644 --- a/web/scanEngine/views.py +++ b/web/scanEngine/views.py @@ -462,15 +462,28 @@ def llm_toolkit_section(request, slug): context = {} list_all_models_url = f'{OLLAMA_INSTANCE}/api/tags' response = requests.get(list_all_models_url) + all_models = [] + all_models = DEFAULT_GPT_MODELS.copy() if response.status_code == 200: models = response.json() - models = models.get('models') + ollama_models = models.get('models') date_format = "%Y-%m-%dT%H:%M:%S" - models = [{**model, 'modified_at': datetime.strptime(model['modified_at'].split('.')[0], date_format)} for model in models] - context['installed_models'] = models - print(context['installed_models']) + for model in ollama_models: + all_models.append({**model, + 'modified_at': datetime.strptime(model['modified_at'].split('.')[0], date_format), + 'is_local': True + }) + # find selected model name from db + selected_model = OllamaSettings.objects.first() + if selected_model: + selected_model = {'selected_model': selected_model.selected_model} else: - context['installed_models'] = [] + # use gpt3.5-turbo as default + selected_model = {'selected_model': 'gpt-3.5-turbo'} + for models in all_models: + if models['name'] == selected_model['selected_model']: + models['selected'] = True + context['installed_models'] = all_models return render(request, 'scanEngine/settings/llm_toolkit.html', context) From 75cadb9ad709f19ce42a40454cc68cb982fbaeaf Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sun, 21 Apr 2024 11:39:49 +0530 Subject: [PATCH 16/18] fix using ollama vs openai --- web/api/views.py | 9 +++++- .../0011_ollamasettings_is_ollama.py | 18 +++++++++++ ...name_is_ollama_ollamasettings_is_openai.py | 18 +++++++++++ .../migrations/0013_auto_20240421_0507.py | 22 +++++++++++++ ...ame_is_ollama_ollamasettings_use_ollama.py | 18 +++++++++++ web/dashboard/models.py | 1 + web/reNgine/gpt.py | 32 +++++++++++++------ web/scanEngine/views.py | 9 +++--- 8 files changed, 113 insertions(+), 14 deletions(-) create mode 100644 web/dashboard/migrations/0011_ollamasettings_is_ollama.py create mode 100644 web/dashboard/migrations/0012_rename_is_ollama_ollamasettings_is_openai.py create mode 100644 web/dashboard/migrations/0013_auto_20240421_0507.py create mode 100644 web/dashboard/migrations/0014_rename_is_ollama_ollamasettings_use_ollama.py diff --git a/web/api/views.py b/web/api/views.py index fe6cd6cf..45fe5df4 100644 --- a/web/api/views.py +++ b/web/api/views.py @@ -88,12 +88,19 @@ class OllamaManager(APIView): def put(self, request): req = self.request model_name = req.query_params.get('model') + # check if model_name is in DEFAULT_GPT_MODELS response = { 'status': False } + use_ollama = True + if any(model['name'] == model_name for model in DEFAULT_GPT_MODELS): + use_ollama = False try: OllamaSettings.objects.update_or_create( - defaults={'selected_model': model_name}, + defaults={ + 'selected_model': model_name, + 'use_ollama': use_ollama + }, id=1 ) response['status'] = True diff --git a/web/dashboard/migrations/0011_ollamasettings_is_ollama.py b/web/dashboard/migrations/0011_ollamasettings_is_ollama.py new file mode 100644 index 00000000..aefdef37 --- /dev/null +++ b/web/dashboard/migrations/0011_ollamasettings_is_ollama.py @@ -0,0 +1,18 @@ +# Generated by Django 3.2.4 on 2024-04-21 05:06 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('dashboard', '0010_ollamasettings'), + ] + + operations = [ + migrations.AddField( + model_name='ollamasettings', + name='is_ollama', + field=models.BooleanField(default=False), + ), + ] diff --git a/web/dashboard/migrations/0012_rename_is_ollama_ollamasettings_is_openai.py b/web/dashboard/migrations/0012_rename_is_ollama_ollamasettings_is_openai.py new file mode 100644 index 00000000..52bdf6ff --- /dev/null +++ b/web/dashboard/migrations/0012_rename_is_ollama_ollamasettings_is_openai.py @@ -0,0 +1,18 @@ +# Generated by Django 3.2.4 on 2024-04-21 05:06 + +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('dashboard', '0011_ollamasettings_is_ollama'), + ] + + operations = [ + migrations.RenameField( + model_name='ollamasettings', + old_name='is_ollama', + new_name='is_openai', + ), + ] diff --git a/web/dashboard/migrations/0013_auto_20240421_0507.py b/web/dashboard/migrations/0013_auto_20240421_0507.py new file mode 100644 index 00000000..11ab6594 --- /dev/null +++ b/web/dashboard/migrations/0013_auto_20240421_0507.py @@ -0,0 +1,22 @@ +# Generated by Django 3.2.4 on 2024-04-21 05:07 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('dashboard', '0012_rename_is_ollama_ollamasettings_is_openai'), + ] + + operations = [ + migrations.RemoveField( + model_name='ollamasettings', + name='is_openai', + ), + migrations.AddField( + model_name='ollamasettings', + name='is_ollama', + field=models.BooleanField(default=True), + ), + ] diff --git a/web/dashboard/migrations/0014_rename_is_ollama_ollamasettings_use_ollama.py b/web/dashboard/migrations/0014_rename_is_ollama_ollamasettings_use_ollama.py new file mode 100644 index 00000000..a201df3a --- /dev/null +++ b/web/dashboard/migrations/0014_rename_is_ollama_ollamasettings_use_ollama.py @@ -0,0 +1,18 @@ +# Generated by Django 3.2.4 on 2024-04-21 05:08 + +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('dashboard', '0013_auto_20240421_0507'), + ] + + operations = [ + migrations.RenameField( + model_name='ollamasettings', + old_name='is_ollama', + new_name='use_ollama', + ), + ] diff --git a/web/dashboard/models.py b/web/dashboard/models.py index 58bb2183..8ed77dd4 100644 --- a/web/dashboard/models.py +++ b/web/dashboard/models.py @@ -29,6 +29,7 @@ class OpenAiAPIKey(models.Model): class OllamaSettings(models.Model): id = models.AutoField(primary_key=True) selected_model = models.CharField(max_length=500) + use_ollama = models.BooleanField(default=True) def __str__(self): return self.selected_model diff --git a/web/reNgine/gpt.py b/web/reNgine/gpt.py index 2e6682ff..389da1fc 100644 --- a/web/reNgine/gpt.py +++ b/web/reNgine/gpt.py @@ -1,16 +1,19 @@ import openai import re from reNgine.common_func import get_open_ai_key, extract_between -from reNgine.definitions import VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE, ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT -from langchain.llms import Ollama +from reNgine.definitions import VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE, ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT, OLLAMA_INSTANCE +from langchain_community.llms import Ollama + +from dashboard.models import OllamaSettings class GPTVulnerabilityReportGenerator: def __init__(self): - self.api_key = get_open_ai_key() - self.model_name = 'gpt-3.5-turbo' - if not self.api_key: - self.ollama = Ollama(base_url='http://ollama:11434', model="llama2-uncensored") + selected_model = OllamaSettings.objects.first() + self.model_name = selected_model.selected_model if selected_model else 'gpt-3.5-turbo' + self.use_ollama = selected_model.use_ollama if selected_model else False + self.openai_api_key = None + self.ollama = None def get_vulnerability_description(self, description): """Generate Vulnerability Description using GPT. @@ -26,12 +29,23 @@ class GPTVulnerabilityReportGenerator: 'references': (list) of urls } """ - if not self.api_key: - prompt = ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT + "\nUser: " + input + print(f"Generating Vulnerability Description for: {description}") + if self.use_ollama: + prompt = VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE + "\nUser: " + description + self.ollama = Ollama( + base_url=OLLAMA_INSTANCE, + model=self.model_name + ) response_content = self.ollama(prompt) else: - openai.api_key = self.api_key + openai_api_key = get_open_ai_key() + if not openai_api_key: + return { + 'status': False, + 'error': 'OpenAI API Key not set' + } try: + openai.api_key = openai_api_key gpt_response = openai.ChatCompletion.create( model=self.model_name, messages=[ diff --git a/web/scanEngine/views.py b/web/scanEngine/views.py index b4f19574..d1eb2ef5 100644 --- a/web/scanEngine/views.py +++ b/web/scanEngine/views.py @@ -463,6 +463,7 @@ def llm_toolkit_section(request, slug): list_all_models_url = f'{OLLAMA_INSTANCE}/api/tags' response = requests.get(list_all_models_url) all_models = [] + selected_model = None all_models = DEFAULT_GPT_MODELS.copy() if response.status_code == 200: models = response.json() @@ -471,7 +472,7 @@ def llm_toolkit_section(request, slug): for model in ollama_models: all_models.append({**model, 'modified_at': datetime.strptime(model['modified_at'].split('.')[0], date_format), - 'is_local': True + 'is_local': True, }) # find selected model name from db selected_model = OllamaSettings.objects.first() @@ -480,9 +481,9 @@ def llm_toolkit_section(request, slug): else: # use gpt3.5-turbo as default selected_model = {'selected_model': 'gpt-3.5-turbo'} - for models in all_models: - if models['name'] == selected_model['selected_model']: - models['selected'] = True + for model in all_models: + if model['name'] == selected_model['selected_model']: + model['selected'] = True context['installed_models'] = all_models return render(request, 'scanEngine/settings/llm_toolkit.html', context) From bbf3ef8aa16602ac1307ba6200b25972c9c80e2f Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sun, 12 May 2024 18:16:57 +0530 Subject: [PATCH 17/18] fix requirements --- web/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/web/requirements.txt b/web/requirements.txt index e402d590..fd548693 100644 --- a/web/requirements.txt +++ b/web/requirements.txt @@ -17,6 +17,7 @@ drf-yasg==1.21.3 gunicorn==22.0.0 gevent==24.2.1 humanize==4.3.0 +langchain==0.1.0 Markdown==3.3.4 metafinder==1.2 netaddr==0.8.0 @@ -37,4 +38,3 @@ whatportis weasyprint==53.3 wafw00f==2.2.0 xmltodict==0.13.0 -langchain==0.0.343 From 95f5384d11dd6962aaebfee9a3eb40823a02c96d Mon Sep 17 00:00:00 2001 From: Yogesh Ojha Date: Sun, 12 May 2024 18:17:15 +0530 Subject: [PATCH 18/18] show error message when openai key is not set --- .../templates/scanEngine/settings/llm_toolkit.html | 13 +++++++++++-- web/scanEngine/views.py | 4 ++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html index b6c87c7e..fa6d304e 100644 --- a/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html +++ b/web/scanEngine/templates/scanEngine/settings/llm_toolkit.html @@ -14,7 +14,7 @@ LLM Toolkit {% endblock breadcrumb_title %} {% block page_title %} -LLM Toolkit +LLM Toolkit (Beta) {% endblock page_title %} {% block main_content %} @@ -22,9 +22,13 @@ LLM Toolkit
    + LLM Toolkit includes the ability to download new LLMs, view available models, and delete models no longer needed, and also choose between various models. +

    reNgine makes use of various LLMs to enhance reporting process. Using various LLM AI Models penetration testers will be able to to generate detailed, insightful penetration testing reports.
    - LLM Toolkit includes the ability to download new LLMs, view available models, and delete models no longer needed, and also choose between various models. + If you are using custom LLM models, it is expected that response time are much slower in CPU. We recommend using GPU for better performance. Models such as llama2, or llama3 requires significant computation and GPU are required. Having only CPU will result in slow response time. +
    + OpenAI GPT models do not run locally, hence the requirement of GPU is not necessary.

    @@ -35,6 +39,11 @@ LLM Toolkit
    {{installed_models|length}} available Models
    + {% if openai_key_error %} + + {% endif %}
    {% for model in installed_models %}
    diff --git a/web/scanEngine/views.py b/web/scanEngine/views.py index d1eb2ef5..0404e8c4 100644 --- a/web/scanEngine/views.py +++ b/web/scanEngine/views.py @@ -485,6 +485,10 @@ def llm_toolkit_section(request, slug): if model['name'] == selected_model['selected_model']: model['selected'] = True context['installed_models'] = all_models + # show error message for openai key, if any gpt is selected + openai_key = get_open_ai_key() + if not openai_key and 'gpt' in selected_model['selected_model']: + context['openai_key_error'] = True return render(request, 'scanEngine/settings/llm_toolkit.html', context)