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
+from llm_handler import send_to_llm, set_local_model_base_url, set_anything_llm_workspace
 from main import generate_data, PROMPT_1
 from topics import TOPICS
 from system_messages import SYSTEM_MESSAGES_VODALUS
@@ -264,8 +264,13 @@ 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):
+def chat_with_llm(message, history, selected_llm, base_url, anything_llm_workspace):
     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]})
@@ -407,7 +412,9 @@ with demo:
 
         with gr.Column(scale=1):
             gr.Markdown("## AI Assistant")
-            selected_llm = gr.Radio(["local-model", "anything-llm"], label="Select LLM", value="local-model")
+            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")
@@ -507,7 +514,7 @@ with demo:
         outputs=[generation_status, generation_output]
     )
 
-    msg.submit(chat_with_llm, [msg, chatbot, selected_llm], [chatbot])
+    msg.submit(chat_with_llm, [msg, chatbot, selected_llm, base_url, anything_llm_workspace], [chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)
 
     # Update chat context when navigating rows or loading dataset
@@ -518,5 +525,11 @@ 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)
+