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