added custom client for ollama

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.

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)