Update app.py

app.py

@@ -1,22 +1,21 @@
-import random
-import datetime
-import sys
 import os
+import sys
 import torch
-import logging
 import json
+import logging
+import gradio as gr
 from importlib.resources import files
 from txagent import TxAgent
 from tooluniverse import ToolUniverse
-import gradio as gr
 
-# Set up logging
+# Logging setup
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 
+# Paths and environment
 current_dir = os.path.dirname(os.path.abspath(__file__))
 os.environ["MKL_THREADING_LAYER"] = "GNU"
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -39,22 +38,12 @@ chat_css = """
 .gr-button svg { width: 32px !important; height: 32px !important; }
 """
 
-def validate_message(message: str) -> bool:
-    """Validate that the message meets minimum requirements."""
-    if not message or not isinstance(message, str):
-        return False
-    # Remove whitespace and check length
-    clean_msg = message.strip()
-    return len(clean_msg) >= 10
-
 def safe_load_embeddings(filepath: str) -> any:
     try:
-        # First try with weights_only=True (secure mode)
         return torch.load(filepath, weights_only=True)
     except Exception as e:
         logger.warning(f"Secure load failed, trying with weights_only=False: {str(e)}")
         try:
-            # If that fails, try with weights_only=False (less secure)
             return torch.load(filepath, weights_only=False)
         except Exception as e:
             logger.error(f"Failed to load embeddings: {str(e)}")
@@ -85,7 +74,6 @@ def patch_embedding_loading():
 
                 if current_count != embedding_count:
                     logger.warning(f"Tool count mismatch (tools: {current_count}, embeddings: {embedding_count})")
-
                     if current_count < embedding_count:
                         self.tool_desc_embedding = self.tool_desc_embedding[:current_count]
                         logger.info(f"Truncated embeddings to match {current_count} tools")
@@ -149,18 +137,22 @@ def create_agent():
         logger.error(f"Failed to create agent: {str(e)}")
         raise
 
-def respond(message, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
-    # Validate the message first
-    if not validate_message(message):
-        error_msg = "Please provide a valid message with a string longer than 10 characters."
-        logger.warning(f"Message validation failed: {message}")
-        return history + [{"role": "user", "content": message}, {"role": "assistant", "content": error_msg}]
+def respond(chat_history, history, temperature, max_new_tokens, max_tokens, multi_agent, conversation, max_round):
+    if not chat_history:
+        return [["assistant", "Please provide a message."]]
+
+    message = chat_history[-1][1] if isinstance(chat_history[-1], (list, tuple)) else chat_history[-1]
+
+    if not isinstance(message, str) or len(message.strip()) <= 10:
+        return chat_history + [["assistant", "Please provide a valid message with a string longer than 10 characters."]]
 
     updated_history = history + [{"role": "user", "content": message}]
-    logger.debug(f"\n==== DEBUG ====\nUser Message: {message}\nFull History: {updated_history}\n================\n")
+    print("\n==== DEBUG ====")
+    print("User Message:", message)
+    print("Full History:", updated_history)
+    print("================\n")
 
     try:
-        # Ensure correct format for run_gradio_chat
         formatted_history = [(m["role"], m["content"]) for m in updated_history]
 
         response_generator = agent.run_gradio_chat(
@@ -173,29 +165,22 @@ def respond(message, history, temperature, max_new_tokens, max_tokens, multi_age
             max_round
         )
     except Exception as e:
-        error_msg = f"Error processing your request: {str(e)}"
-        logger.error(f"Error in respond function: {str(e)}")
-        return history + [{"role": "user", "content": message}, {"role": "assistant", "content": error_msg}]
-
-    collected = ""
-    try:
+        updated_history.append({"role": "assistant", "content": f"Error: {str(e)}"})
+    else:
+        collected = ""
         for chunk in response_generator:
             if isinstance(chunk, dict):
                 collected += chunk.get("content", "")
             else:
                 collected += str(chunk)
-    except Exception as e:
-        error_msg = f"Error generating response: {str(e)}"
-        logger.error(f"Error in response generation: {str(e)}")
-        return history + [{"role": "user", "content": message}, {"role": "assistant", "content": error_msg}]
+        updated_history.append({"role": "assistant", "content": collected})
 
-    return history + [{"role": "user", "content": message}, {"role": "assistant", "content": collected}]
+    # Return formatted history to Gradio
+    return [(m["role"], m["content"]) for m in updated_history]
 
 def create_demo(agent):
     with gr.Blocks(css=chat_css) as demo:
         chatbot = gr.Chatbot(label="TxAgent", type="messages")
-        with gr.Row():
-            msg = gr.Textbox(label="Your question", placeholder="Enter your biomedical question here (minimum 10 characters)...")
         with gr.Row():
             temp = gr.Slider(0, 1, value=0.3, label="Temperature")
             max_new_tokens = gr.Slider(128, 4096, value=1024, label="Max New Tokens")
@@ -207,9 +192,10 @@ def create_demo(agent):
 
         submit.click(
             respond,
-            inputs=[msg, chatbot, temp, max_new_tokens, max_tokens, multi_agent, gr.State([]), max_rounds],
+            inputs=[chatbot, chatbot, temp, max_new_tokens, max_tokens, multi_agent, gr.State([]), max_rounds],
             outputs=[chatbot]
         )
+
     return demo
 
 def main():
@@ -217,10 +203,10 @@ def main():
         global agent
         agent = create_agent()
         demo = create_demo(agent)
-        demo.launch(server_name="0.0.0.0", server_port=7860)
+        demo.launch()
     except Exception as e:
         logger.error(f"Application failed to start: {str(e)}")
         raise
 
 if __name__ == "__main__":
-    main()
+    main()