updates

.gitignore

@@ -3,3 +3,4 @@ translations/*
 /server/__pycache__/*
 *.pyc
 server/__pycache__/backend.cpython-310.pyc
+*.mo

server/backend.py

@@ -4,7 +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 server.utils import check_model, make_prompt , make_simple_prompt
 from langchain_community.llms import Ollama
 import requests
 
@@ -60,9 +60,11 @@ class Backend_Api:
         try:
             api_key = request.json['api_key']
             jailbreak = request.json['jailbreak']
+            print("jailbreak:",jailbreak)
             model = request.json['model']
             check_model(model)
             messages = build_messages(jailbreak)
+            print("messages:",messages)
             local_mode_1=True
             local_model_2 =False
             print(model)
@@ -75,9 +77,16 @@ class Backend_Api:
 
             if local_mode_1:
                 content=messages[0]['content']
+                input = request.json['meta']['content']['parts'][0]['content']
+                #prompt=make_simple_prompt(input,messages)
+                prompt=make_prompt(input,messages,model)
                 llm = Ollama(model=model)
+                print("messages",messages)
+                print("len(messages)",len(messages))
                 print("content:",content)
-                response = llm.invoke(content)
+                print("input",input)
+                print("prompt",prompt)
+                response = llm.invoke(prompt)
                 return response                        
             elif local_model_2:
                 # Use the local model to generate the response

server/utils.py

@@ -50,4 +50,80 @@ def check_model(model_name):
                 print(f"Failed to download model '{model_name}': {e}")
                 return
     else:
-        print("OK")
+        print("OK")
+
+
+
+def make_simple_prompt(input, messages):
+    """
+    Create a simple prompt based on the input and messages.
+    
+    :param input: str, input message from the user
+    :param messages: list, conversation history as a list of dictionaries containing 'role' and 'content'
+    :return: str, generated prompt
+    """
+    if len(messages) == 1:
+        prompt = f'''You are a friendly AI companion.
+You should answer what the user request.
+user: {input}'''
+    else:
+        conversation_history = '\n'.join(
+            f"{message['role']}: {message['content']}" for message in reversed(messages[:-1])
+        )
+        prompt = f'''You are a friendly AI companion.
+history: {conversation_history}.
+You should answer what the user request.
+user: {input}'''
+
+    print(prompt)
+    return prompt
+
+
+def make_prompt(input, messages, model):
+    """
+    Create a prompt based on the input, messages, and model used.
+    
+    :param input: str, input message from the user
+    :param messages: list, conversation history as a list of dictionaries containing 'role' and 'content'
+    :param model: str, name of the model ("llama3", "mistral", or other)
+    :return: str, generated prompt
+    """
+    if model == "llama3":
+        # Special Tokens used with Meta Llama 3
+        BEGIN_OF_TEXT = "<|begin_of_text|>"
+        EOT_ID = "<|eot_id|>"
+        START_HEADER_ID = "<|start_header_id|>"
+        END_HEADER_ID = "<|end_header_id|>"
+    elif model == "mistral":
+        # Special tokens Mistral
+        BEGIN_OF_TEXT = "<s>"
+        EOT_ID = "</s>"
+        START_HEADER_ID = ""  # Not applicable to Mistral
+        END_HEADER_ID = ""  # Not applicable to Mistral
+    else:
+        # No Special tokens
+        BEGIN_OF_TEXT = ""
+        EOT_ID = ""
+        START_HEADER_ID = ""
+        END_HEADER_ID = ""
+
+    if len(messages) == 1:
+        prompt = f'''{BEGIN_OF_TEXT}{START_HEADER_ID}system{END_HEADER_ID}
+You are a friendly AI companion.
+{EOT_ID}{START_HEADER_ID}user{END_HEADER_ID}
+{input}
+{EOT_ID}'''
+    else:
+        conversation_history = '\n'.join(
+            f"{START_HEADER_ID}{message['role']}{END_HEADER_ID}\n{message['content']}{EOT_ID}" for message in reversed(messages[:-1])
+        )
+        prompt = f'''{BEGIN_OF_TEXT}{START_HEADER_ID}system{END_HEADER_ID}
+You are a friendly AI companion.
+history:
+{conversation_history}
+{EOT_ID}{START_HEADER_ID}user{END_HEADER_ID}
+{input}
+{EOT_ID}'''
+
+    print(prompt)
+    return prompt