Update app.py

app.py

@@ -13,9 +13,9 @@ import json
 now = datetime.now()
 date_time_str = now.strftime("%Y-%m-%d %H:%M:%S")
 
-client = InferenceClient(
-    "mistralai/Mixtral-8x7B-Instruct-v0.1"
-)
+# Define the model globally (or pass it as an argument to main)
+model = "mistralai/Mixtral-8x7B-Instruct-v0.1" 
+client = InferenceClient(model)
 
 # --- Set up logging ---
 logging.basicConfig(
@@ -275,8 +275,7 @@ def project_explorer(path):
     tree = get_file_tree(path)
     display_file_tree(tree)
 
-def chat_app_logic(message, history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty):
-    # Your existing code here
+def chat_app_logic(message, history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty, model):  # Add 'model' as an argument
     
     try:
         # Attempt to join the generator output
@@ -302,6 +301,41 @@ def chat_app_logic(message, history, purpose, agent_name, sys_prompt, temperatur
             elif isinstance(part, dict) and 'content' in part:
                 response_parts.append(part['content']),
         
+        response = ''.join(response_parts,
+        # Run the model and get the response (convert generator to string)
+            prompt=message,
+            history=history,
+            agent_name=agent_name,
+            sys_prompt=sys_prompt,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+                          )
+    try:
+        # Attempt to join the generator output
+        response = ''.join(generate(
+            model=model,  # Now you can use 'model' here
+            messages=messages,
+            stream=True,
+            temperature=0.7,
+            max_tokens=1500
+        ))
+    except TypeError:
+        # If joining fails, collect the output in a list
+        response_parts = []
+        for part in generate(
+            model=model,  # Now you can use 'model' here
+            messages=messages,
+            stream=True,
+            temperature=0.7,
+            max_tokens=1500
+        ):
+            if isinstance(part, str):
+                response_parts.append(part)
+            elif isinstance(part, dict) and 'content' in part:
+                response_parts.append(part['content']),
+        
         response = ''.join(response_parts,
         # Run the model and get the response (convert generator to string)
             prompt=message,
@@ -319,6 +353,7 @@ def chat_app_logic(message, history, purpose, agent_name, sys_prompt, temperatur
     return history
 
 def main():
+    
     with gr.Blocks() as demo:
         gr.Markdown("## FragMixt")
         gr.Markdown("### Agents w/ Agents")
@@ -346,15 +381,16 @@ def main():
             explore_button = gr.Button(value="Explore")
             project_output = gr.Textbox(label="File Tree", lines=20)
 
-        # Chat App Logic Tab
+                # Chat App Logic Tab
         with gr.Tab("Chat App"):
             history = gr.State([])
             for example in examples:
-                gr.Button(value=example[0]).click(lambda: chat_app_logic(example[0], history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty), outputs=chatbot)
+                gr.Button(value=example[0]).click(lambda: chat_app_logic(example[0], history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty, model), outputs=chatbot)
 
             # Connect components to the chat app logic
-            submit_button.click(chat_app_logic, inputs=[message, history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty], outputs=chatbot)
-            message.submit(chat_app_logic, inputs=[message, history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty], outputs=chatbot)
+            submit_button.click(chat_app_logic, inputs=[message, history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty, model], outputs=chatbot)  # Pass 'model'
+            message.submit(chat_app_logic, inputs=[message, history, purpose, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty, model], outputs=chatbot)  # Pass 'model'
+            
 
         # Connect components to the project explorer
         explore_button.click(project_explorer, inputs=project_path, outputs=project_output)