Adding models

Dockerfile

@@ -34,7 +34,7 @@ ENV PATH="/app/venv/bin:/root/.ollama/bin:$PATH"
 # Expose the server port
 EXPOSE 7860
 EXPOSE 11434
-
+EXPOSE 1338
 # Copy the entry point script
 COPY entrypoint.sh /entrypoint.sh
 RUN chmod +x /entrypoint.sh

server/backend.py

@@ -4,6 +4,7 @@ from g4f import ChatCompletion
 from flask import request, Response, stream_with_context
 from requests import get
 from server.config import special_instructions
+from server.utils import check_model
 from langchain_community.llms import Ollama
 import requests
 
@@ -60,7 +61,7 @@ class Backend_Api:
             api_key = request.json['api_key']
             jailbreak = request.json['jailbreak']
             model = request.json['model']
-            
+            check_model(model)
             messages = build_messages(jailbreak)
             local_mode_1=True
             local_model_2 =False

server/utils.py

@@ -0,0 +1,53 @@
+import subprocess
+def check_model_exists(model_name):
+    try:
+        # List available models
+        output = subprocess.check_output("ollama list", shell=True, stderr=subprocess.STDOUT, universal_newlines=True)
+        available_models = [line.split()[0] for line in output.strip().split('\n')[1:]]
+        return any(model_name in model for model in available_models)
+    except subprocess.CalledProcessError as e:
+        print(f"Error checking models: {e.output}")
+        return False
+    except Exception as e:
+        print(f"An unexpected error occurred: {str(e)}")
+        return False
+    
+    
+def download_model(model_name):
+    remote_models=['llama3',
+    'llama3:70b',
+    'phi3',
+    'mistral',
+    'neural-chat',
+    'starling-lm',
+    'codellama',
+    'llama2-uncensored',
+    'llava',
+    'gemma:2b',
+    'gemma:7b',
+    'solar']
+    if model_name in remote_models:
+        try:
+            # Download the model
+            print(f"Downloading model '{model_name}'...")
+            subprocess.check_call(f"ollama pull {model_name}", shell=True)
+            print(f"Model '{model_name}' downloaded successfully.")
+        except subprocess.CalledProcessError as e:
+            print(f"Error downloading model: {e.output}")
+            raise e
+        except Exception as e:
+            print(f"An unexpected error occurred: {str(e)}")
+            raise e
+    else:
+        print("Not supported model currently")
+
+
+def check_model(model_name):
+    if not check_model_exists(model_name):
+            try:
+                download_model(model_name)
+            except Exception as e:
+                print(f"Failed to download model '{model_name}': {e}")
+                return
+    else:
+        print("OK")