Update app.py

app.py

@@ -5,7 +5,7 @@ from datetime import datetime
 from typing import Literal
 import os
 import importlib
-from llm_handler import send_to_llm, set_local_model_base_url, set_anything_llm_workspace
+from llm_handler import send_to_llm
 from main import generate_data, PROMPT_1
 from topics import TOPICS
 from system_messages import SYSTEM_MESSAGES_VODALUS
@@ -264,20 +264,15 @@ def save_dataset_config(system_messages, prompt_1, topics):
     return "Dataset configuration saved successfully"
 
 # Modify the chat_with_llm function to use Gradio's built-in async capabilities
-def chat_with_llm(message, history, selected_llm, base_url, anything_llm_workspace):
+def chat_with_llm(message, history):
     try:
-        if selected_llm == "local-model":
-            set_local_model_base_url(base_url)
-        elif selected_llm == "anything-llm":
-            set_anything_llm_workspace(anything_llm_workspace)
-        
         msg_list = [{"role": "system", "content": "You are an AI assistant helping with dataset annotation and quality checking."}]
         for h in history:
             msg_list.append({"role": "user", "content": h[0]})
             msg_list.append({"role": "assistant", "content": h[1]})
         msg_list.append({"role": "user", "content": message})
 
-        response, _ = send_to_llm(selected_llm, msg_list)
+        response, _ = send_to_llm("llamanet", msg_list)
         
         return history + [[message, response]]
     except Exception as e:
@@ -297,12 +292,12 @@ def update_chat_context(row_data, index, total, quality, high_quality_tags, low_
     return [[None, context]]  # Return as a list of message pairs
 
 # Add this function to handle dataset generation
-async def run_generate_dataset(num_workers, num_generations, output_file_path, selected_llm):
+async def run_generate_dataset(num_workers, num_generations, output_file_path):
     generated_data = []
     for _ in range(num_generations):
         topic_selected = random.choice(TOPICS)
         system_message_selected = random.choice(SYSTEM_MESSAGES_VODALUS)
-        data = await generate_data(topic_selected, PROMPT_1, system_message_selected, output_file_path, selected_llm)
+        data = await generate_data(topic_selected, PROMPT_1, system_message_selected, output_file_path)
         if data:
             generated_data.append(json.dumps(data))
     
@@ -412,9 +407,6 @@ with demo:
 
         with gr.Column(scale=1):
             gr.Markdown("## AI Assistant")
-            selected_llm = gr.Radio(["local-model", "anything-llm", "llamacpp"], label="Select LLM", value="local-model")
-            base_url = gr.Textbox(label="Base URL for local-model", value="http://localhost:11434/v1", visible=False)
-            anything_llm_workspace = gr.Textbox(label="AnythingLLM Workspace", value="<input-workspace-name-here>", visible=False)
             chatbot = gr.Chatbot(height=600)
             msg = gr.Textbox(label="Chat with AI Assistant")
             clear = gr.Button("Clear")
@@ -510,11 +502,11 @@ with demo:
 
     start_generation_btn.click(
         run_generate_dataset,
-        inputs=[num_workers, num_generations, output_file_path, selected_llm],
+        inputs=[num_workers, num_generations, output_file_path],
         outputs=[generation_status, generation_output]
     )
 
-    msg.submit(chat_with_llm, [msg, chatbot, selected_llm, base_url, anything_llm_workspace], [chatbot])
+    msg.submit(chat_with_llm, [msg, chatbot], [chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)
 
     # Update chat context when navigating rows or loading dataset
@@ -525,11 +517,6 @@ with demo:
             outputs=[chatbot]
         )
 
-    def toggle_input_visibility(llm):
-        return gr.update(visible=llm == "local-model"), gr.update(visible=llm == "anything-llm")
-
-    selected_llm.change(toggle_input_visibility, inputs=[selected_llm], outputs=[base_url, anything_llm_workspace])
-
 if __name__ == "__main__":
     demo.launch(share=True)