Update app.py

app.py

@@ -9,6 +9,7 @@ import gradio as gr
 from huggingface_hub import InferenceClient
 from safe_search import safe_search
 from i_search import google, i_search as i_s
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # --- Configuration ---
 VERBOSE = True
@@ -141,6 +142,31 @@ def generate(
         logging.info(LOG_RESPONSE.format(resp=response))
     return response.text
 
+# --- Mixtral Integration ---
+def mixtral_generate(
+    prompt: str,
+    history: List[Tuple[str, str]],
+    agent_name: str = agents[0],
+    sys_prompt: str = "",
+    temperature: float = TEMPERATURE,
+    max_new_tokens: int = MAX_TOKENS,
+    top_p: float = TOP_P,
+    repetition_penalty: float = REPETITION_PENALTY,
+) -> str:
+    """Generates a response using the Mixtral model."""
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
+    content = PREFIX.format(
+        date_time_str=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        purpose=f"Generating response as {agent_name}",
+        safe_search=safe_search,
+    ) + sys_prompt + "\n" + prompt
+
+    inputs = tokenizer(content, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
 def main():
     """Main function to launch the Gradio interface."""
     with gr.Blocks() as demo:
@@ -242,7 +268,8 @@ def main():
         ) -> Tuple[List[Tuple[str, str]], List[Tuple[str, str]]]:
             """Handles the chat interaction, generating responses and updating history."""
             prompt = format_prompt(message, history)
-            response = generate(
+            # Use Mixtral for generation
+            response = mixtral_generate(
                 prompt,
                 history,
                 agent_name,