WEBDIENSTE/rengine
1
0
Fork 0
mirror of https://github.com/yogeshojha/rengine.git synced 2026-09-26 19:54:53 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1259 from yogeshojha/feature/llm

Browse Source 
Add LLM Toolkit
This commit is contained in:
Yogesh Ojha
2024-05-12 19:00:24 +05:30
committed by GitHub
parent 7bb9d42e5c 95f5384d11
commit 56da132afd
20 changed files with 680 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CONTRIBUTORS.md
+1
View File
@@ -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!*
Makefile
+1 -1
View File
@@ -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
# --------------------------
docker-compose.dev.yml
+12
View File
@@ -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:
docker-compose.yml
+10 -1
View File
@@ -143,7 +143,15 @@ 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
networks:
rengine_network:
@@ -157,6 +165,7 @@ volumes:
wordlist:
scan_results:
static_volume:
ollama_data:
secrets:
proxy.ca:
web/api/urls.py
+4
View File
@@ -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(),
web/api/views.py
+76
View File
@@ -33,6 +33,82 @@ 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)
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,
'use_ollama': use_ollama
},
id=1
)
response['status'] = True
except Exception as e:
response['error'] = str(e)
return Response(response)
class GPTAttackSuggestion(APIView):
def get(self, request):
req = self.request
web/dashboard/migrations/0010_ollamasettings.py
+20
View File
@@ -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)),
],
),
]
web/dashboard/migrations/0011_ollamasettings_is_ollama.py
+18
View File
@@ -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),
),
]
web/dashboard/migrations/0012_rename_is_ollama_ollamasettings_is_openai.py
+18
View File
@@ -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',
),
]
web/dashboard/migrations/0013_auto_20240421_0507.py
+22
View File
@@ -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),
),
]
web/dashboard/migrations/0014_rename_is_ollama_ollamasettings_use_ollama.py
+18
View File
@@ -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',
),
]
web/dashboard/models.py
+9
View File
@@ -24,6 +24,15 @@ 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)
use_ollama = models.BooleanField(default=True)
def __str__(self):
return self.selected_model
class NetlasAPIKey(models.Model):
web/reNgine/definitions.py
+46
View File
@@ -436,6 +436,52 @@ PERM_INITATE_SCANS_SUBSCANS = 'initiate_scans_subscans'
FOUR_OH_FOUR_URL = '/404/'
###############################################################################
# OLLAMA DEFINITIONS
###############################################################################
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 = """
You are a highly skilled penetration tester who has recently completed a penetration testing.
web/reNgine/gpt.py
+93 -78
View File
@@ -1,14 +1,20 @@
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 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'
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.
@@ -23,94 +29,103 @@ class GPTVulnerabilityReportGenerator:
'references': (list) of urls
}
"""
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}
]
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 = 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=[
{'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://ollama: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
}
web/requirements.txt
+1
View File
@@ -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
web/scanEngine/templates/scanEngine/settings/api.html
+4 -4
View File
@@ -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
<div class="col-12">
<div class="p-sm-3">
<div class="mb-3">
<label for="key_openai" class="form-label">OpenAI <span class="ms-1 badge bg-soft-danger text-danger">🔥 Recommended</span><span class="ms-1 badge bg-soft-primary text-primary">Experimental</span></label>
<p class="text-muted">OpenAI keys will be used to generate vulnerability description, remediation, impact and vulnerability report writing using ChatGPT.</p>
<label for="key_openai" class="form-label">OpenAI</label>
<p class="text-muted">OpenAI keys will be used to generate vulnerability description, remediation, impact and vulnerability report writing using GPT.</p>
{% if openai_key %}
<input class="form-control" type="text" id="key_openai" name="key_openai" placeholder="Enter OpenAI Key" value="{{openai_key}}">
{% else %}
@@ -38,7 +38,7 @@ API Vault
</div>
<div class="mb-3">
<label for="key_netlas" class="form-label">Netlas</label>
<p class="text-muted">Netlas keys will be used to get whois information and other OSINT data.</p>
<p class="text-muted">Netlas keys will be used to get whois information and other OSINT related data.</p>
{% if netlas_key %}
<input class="form-control" type="text" id="key_netlas" name="key_netlas" placeholder="Enter Netlas Key" value="{{netlas_key}}">
{% else %}
web/scanEngine/templates/scanEngine/settings/llm_toolkit.html
+286
View File
@@ -0,0 +1,286 @@
{% extends 'base/base.html' %}
{% load static %}
{% load humanize %}
{% block title %}
LLM Toolkit
{% endblock title %}
{% block custom_js_css_link %}
{% endblock custom_js_css_link %}
{% block breadcrumb_title %}
<li class="breadcrumb-item"><a href="#">Settings</a></li>
<li class="breadcrumb-item active">LLM Toolkit</li>
{% endblock breadcrumb_title %}
{% block page_title %}
LLM Toolkit (Beta)
{% endblock page_title %}
{% block main_content %}
<div class="row">
<div class="col-12">
<div class="card">
<div class="card-body">
<b>LLM Toolkit</b> includes the ability to download new LLMs, view available models, and delete models no longer needed, and also choose between various models.
<br>
<p>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.
<br>
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. <b>Having only CPU will result in slow response time.</b>
<br>
<b>OpenAI GPT models do not run locally, hence the requirement of GPU is not necessary.</b>
</p>
</div>
</div>
</div>
<div class="row mb-2">
<div class="col-sm-4">
<a href="#" class="btn btn-primary rounded-pill waves-effect waves-light mb-3" onclick="showAddNewModelModal()"><i class="mdi mdi-plus"></i> Add new model</a>
</div>
</div>
<h5>{{installed_models|length}} available Models</h5>
{% if openai_key_error %}
<div class="alert alert-danger border-0 mb-3 mt-3" role="alert">
<b>Warning:</b> GPT model is currently selected and requires API key to be set. Please set the API key in the <a href="/scanEngine/{{current_project.slug}}/api_vault" target="_blank"> API Vault.</a>
</div>
{% endif %}
<div class="row mt-2">
{% for model in installed_models %}
<div class="col-lg-4">
<div class="card project-box">
<div class="card-body">
<div class="dropdown float-end">
<a href="#" class="dropdown-toggle card-drop arrow-none" data-bs-toggle="dropdown" aria-expanded="false">
<i class="mdi mdi-dots-horizontal m-0 text-muted h3"></i>
</a>
<div class="dropdown-menu dropdown-menu-end">
{% if model.is_local %}
<a class="dropdown-item" href="#" onClick=deleteModel('{{model.name}}')>Delete</a>
{% endif %}
{% if not model.selected %}
<a class="dropdown-item" href="#" onClick=selectModel('{{model.name}}')>Use Model</a>
{% endif %}
</div>
</div>
<h4 class="mt-0">
<span class="{% if model.selected %}text-success{% endif %}">{{model.name}} {% if model.selected %}<span class="badge bg-soft-primary text-primary ms-4">Selected Model</span>{% endif %}</span>
</h4>
</p>
<p class="mb-1">
<span class="pe-2 text-nowrap mb-2 d-inline-block">
<i class="mdi mdi-calendar-range text-primary"></i>
Modified <b>{% if model.modified_at %}{{model.modified_at|naturaltime}} {% else %} NA{% endif %}</b>
</span>
<br>
<span class="pe-2 text-nowrap mb-2 d-inline-block">
<i class="mdi mdi-database text-info"></i>
{% if model.is_local %}
Locally installed model
{% else %}
Open AI Model
{% endif %}
</span>
<br>
<span class="pe-2 text-nowrap mb-2 d-inline-block">
<i class="mdi mdi-numeric text-info"></i>
<b>{{model.details.parameter_size}}</b> Parameters
</span>
<span class="text-nowrap mb-2 d-inline-block">
<i class="mdi mdi-family-tree text-success"></i>
<b>{{model.details.family}}</b> Family
</span>
</p>
</div>
</div>
</div>
{% empty %}
<div class="alert alert-danger border-0" role="alert">
No LLM Models are installed. You can install models using the 'Add New LLM' button.
</div>
{% endfor %}
</div>
</div>
{% endblock main_content %}
{% block page_level_script %}
<script type="text/javascript">
function deleteModel(model_name) {
// split model name by : and only use first part
model_name = model_name.split(':')[0];
var url = "/api/tool/ollama/?model=" + model_name;
swal.queue([{
title: 'Are you sure you want to delete this model?',
text: "This action can not be undone.",
icon: 'warning',
showCancelButton: true,
confirmButtonText: 'Delete',
padding: '2em',
showLoaderOnConfirm: true,
preConfirm: function() {
return fetch(url, {
method: 'DELETE',
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 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',
title: 'Oops! Unable to delete the model!'
})
})
}
}])
}
function showAddNewModelModal(){
$('#modal_title').html('Add new LLM Model');
$('#modal-content').empty();
$('#modal-content').append(`
<p>You can find the list of supported models in <a href="https://ollama.com/library" target="_blank">Ollama Library</a></p>
<p>We recommend using llama2-uncensored model for better results.</p>
<div class="mb-3">
<label for="model_name" class="form-label">Model name</label>
<input class="form-control" type="text" id="model_name" required="" placeholder="llama2">
</div>
<div class="mb-3 text-center">
<button class="btn btn-primary float-end" type="submit" onclick="download_model()">Download Model</button>
</div>
`);
$('#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!'
})
})
}
}])
}
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!'
})
})
}
}])
}
</script>
{% endblock page_level_script %}
web/scanEngine/urls.py
+4
View File
@@ -56,6 +56,10 @@ urlpatterns = [
'<slug:slug>/tool_arsenal',
views.tool_arsenal_section,
name='tool_arsenal'),
path(
'<slug:slug>/llm_toolkit',
views.llm_toolkit_section,
name='llm_toolkit'),
path(
'<slug:slug>/rengine_settings',
views.rengine_settings,
web/scanEngine/views.py
+36
View File
@@ -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
@@ -456,6 +457,41 @@ 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 = {}
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()
ollama_models = models.get('models')
date_format = "%Y-%m-%dT%H:%M:%S"
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:
# use gpt3.5-turbo as default
selected_model = {'selected_model': 'gpt-3.5-turbo'}
for model in all_models:
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)
@has_permission_decorator(PERM_MODIFY_SYSTEM_CONFIGURATIONS, redirect_url=FOUR_OH_FOUR_URL)
def api_vault(request, slug):
context = {}
web/templates/base/_items/top_nav.html
+1
View File
@@ -81,6 +81,7 @@
{% endif %}
{% if user|can:'modify_system_configurations' %}
<a href="{% url 'api_vault' current_project.slug %}" class="dropdown-item">API Vault</a>
<a href="{% url 'llm_toolkit' current_project.slug %}" class="dropdown-item">LLM Toolkit</a>
<a href="{% url 'tool_arsenal' current_project.slug %}" class="dropdown-item">Tools Arsenal</a>
{% endif %}
{% if user|can:'modify_scan_report' %}