From 36967f8fc3d9e0da26de3d2f034abafd5cc5cd23 Mon Sep 17 00:00:00 2001
From: Yogesh Ojha <yogesh.ojha11@gmail.com>
Date: Sat, 20 Apr 2024 16:18:59 +0530
Subject: [PATCH] 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