Update app.py

app.py

@@ -10,6 +10,8 @@ from selenium.webdriver.chrome.service import Service
 from selenium.webdriver.chrome.options import Options
 from webdriver_manager.chrome import ChromeDriverManager
 from huggingface_hub import InferenceClient
+import random
+import yaml
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -98,9 +100,70 @@ def handle_input(storage_location, url1, url2, scrape_interval, content_type):
     monitor_urls(storage_location, url1, url2, scrape_interval, content_type)
     return TASK_PROMPT.format(task=current_task, history="\n".join(history))
 
-# Define the chat response function
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# Load custom prompts
+try:
+    with open('custom_prompts.yaml', 'r') as fp:
+        custom_prompts = yaml.safe_load(fp)
+except FileNotFoundError:
+    custom_prompts = {
+        "WEB_DEV": "",
+        "AI_SYSTEM_PROMPT": "",
+        "PYTHON_CODE_DEV": "",
+        "CODE_GENERATION": "",
+        "CODE_INTERPRETATION": "",
+        "CODE_TRANSLATION": "",
+        "CODE_IMPLEMENTATION": ""
+    }
+
+# Define the Mistral inference client
+client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+
+VERBOSE = True
+MAX_HISTORY = 125
+
+def format_prompt(message, history):
+    prompt = "<s>"
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    prompt += f"[INST] {message} [/INST]"
+    return prompt
+
+agents = [
+    "WEB_DEV",
+    "AI_SYSTEM_PROMPT",
+    "PYTHON_CODE_DEV",
+    "CODE_GENERATION",
+    "CODE_INTERPRETATION",
+    "CODE_TRANSLATION",
+    "CODE_IMPLEMENTATION"
+]
+
+def generate(
+        prompt, history, agent_name=agents[0], sys_prompt="", temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.7,
+):
+    seed = random.randint(1, 1111111111111111)
+    agent = custom_prompts[agent_name]
+    
+    system_prompt = agent if sys_prompt == "" else sys_prompt
+    temperature = max(float(temperature), 1e-2)
+    top_p = float(top_p)
+
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=seed,
+    )
+
+    formatted_prompt = format_prompt(f"{system_prompt}\n\n{prompt}", history)
+    output = client.text_generation(formatted_prompt, **generate_kwargs, stream=False, return_full_text=False)
+    
+    return output
 
+# Define the chat response function
 def respond(
     message,
     history,
@@ -109,29 +172,15 @@ def respond(
     temperature,
     top_p,
 ):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
+    response = generate(
+        prompt=message,
+        history=history,
+        sys_prompt=system_message,
         temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
+        max_new_tokens=max_tokens,
+        top_p=top_p
+    )
+    return response
 
 # Create Gradio interface
 def chat_interface(message, system_message, max_tokens, temperature, top_p, storage_location, url1, url2, scrape_interval, content_type):