Update app.py

app.py

@@ -2,10 +2,10 @@ import os
 import json
 import logging
 import torch
-from txagent import TxAgent
 import gradio as gr
 from tooluniverse import ToolUniverse
 import warnings
+from typing import List, Dict, Any
 
 # Suppress specific warnings
 warnings.filterwarnings("ignore", category=UserWarning)
@@ -28,6 +28,7 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 def prepare_tool_files():
+    """Ensure tool files exist and are populated"""
     os.makedirs("./data", exist_ok=True)
     if not os.path.exists(CONFIG["tool_files"]["new_tool"]):
         logger.info("Generating tool list using ToolUniverse...")
@@ -37,7 +38,7 @@ def prepare_tool_files():
             json.dump(tools, f, indent=2)
         logger.info(f"Saved {len(tools)} tools to {CONFIG['tool_files']['new_tool']}")
 
-def safe_load_embeddings(filepath):
+def safe_load_embeddings(filepath: str) -> Any:
     """Safely load embeddings with proper weights_only handling"""
     try:
         # First try with weights_only=True (secure mode)
@@ -54,7 +55,7 @@ def patch_embedding_loading():
         
         original_load = ToolRAGModel.load_tool_desc_embedding
         
-        def patched_load(self, tooluniverse):
+        def patched_load(self, tooluniverse: ToolUniverse) -> bool:
             try:
                 if not os.path.exists(CONFIG["embedding_filename"]):
                     logger.error(f"Embedding file not found: {CONFIG['embedding_filename']}")
@@ -64,7 +65,7 @@ def patch_embedding_loading():
                 self.tool_desc_embedding = safe_load_embeddings(CONFIG["embedding_filename"])
                 
                 # Handle tool count mismatch
-                tools = tooluniverse.tools  # Changed from get_all_tools()
+                tools = tooluniverse.get_all_tools()  # Use get_all_tools() instead of direct access
                 current_count = len(tools)
                 embedding_count = len(self.tool_desc_embedding)
                 
@@ -101,7 +102,8 @@ class TxAgentApp:
         self.agent = None
         self.is_initialized = False
 
-    def initialize(self):
+    def initialize(self) -> str:
+        """Initialize the TxAgent with all required components"""
         if self.is_initialized:
             return "✅ Already initialized"
         
@@ -111,8 +113,8 @@ class TxAgentApp:
             
             logger.info("Initializing TxAgent...")
             self.agent = TxAgent(
-                CONFIG["model_name"],
-                CONFIG["rag_model_name"],
+                model_name=CONFIG["model_name"],
+                rag_model_name=CONFIG["rag_model_name"],
                 tool_files_dict=CONFIG["tool_files"],
                 force_finish=True,
                 enable_checker=True,
@@ -131,36 +133,51 @@ class TxAgentApp:
             logger.error(f"Initialization failed: {str(e)}")
             return f"❌ Initialization failed: {str(e)}"
 
-    def chat(self, message, history):
+    def chat(self, message: str, history: List[List[str]]) -> List[List[str]]:
+        """
+        Handle chat interactions with the TxAgent
+        
+        Args:
+            message: User input message
+            history: Chat history in format [[user_msg, bot_msg], ...]
+            
+        Returns:
+            Updated chat history
+        """
         if not self.is_initialized:
-            return {"role": "assistant", "content": "⚠️ Please initialize the model first"}
+            return history + [["", "⚠️ Please initialize the model first"]]
         
         try:
+            # Convert history to the format TxAgent expects
+            tx_history = []
+            for user_msg, bot_msg in history:
+                tx_history.append({"role": "user", "content": user_msg})
+                if bot_msg:  # Only add bot response if it exists
+                    tx_history.append({"role": "assistant", "content": bot_msg})
+            
+            # Generate response
             response = ""
-            # Modified to use the correct parameter name (max_length instead of max_tokens)
             for chunk in self.agent.run_gradio_chat(
                 message=message,
-                history=history,
+                history=tx_history,
                 temperature=0.3,
                 max_new_tokens=1024,
-                max_length=8192,  # Changed from max_tokens
-                multi_agent=False,
-                conversation=[],
+                max_token=8192,  # Note: Using max_token instead of max_length
+                call_agent=False,
+                conversation=None,
                 max_round=30
             ):
-                response += chunk
+                response = chunk  # Get the final response
             
-            # Format response in the expected messages format
-            return [
-                {"role": "user", "content": message},
-                {"role": "assistant", "content": response}
-            ]
+            # Format response for Gradio Chatbot
+            return history + [[message, response]]
                 
         except Exception as e:
             logger.error(f"Chat error: {str(e)}")
-            return {"role": "assistant", "content": f"Error: {str(e)}"}
+            return history + [["", f"Error: {str(e)}"]]
 
-def create_interface():
+def create_interface() -> gr.Blocks:
+    """Create the Gradio interface"""
     app = TxAgentApp()
     
     with gr.Blocks(
@@ -195,25 +212,23 @@ def create_interface():
             inputs=msg
         )
         
-        def wrapper_initialize():
+        def wrapper_initialize() -> tuple:
+            """Wrapper for initialization with UI updates"""
             status = app.initialize()
             return status, gr.update(interactive=False)
         
-        def wrapper_chat(message, chat_history):
-            response = app.chat(message, chat_history)
-            if isinstance(response, dict):  # Error case
-                return chat_history + [response]
-            return response  # Normal case
-        
         init_btn.click(
             fn=wrapper_initialize,
             outputs=[init_status, init_btn]
         )
         
         msg.submit(
-            fn=wrapper_chat,
+            fn=app.chat,
             inputs=[msg, chatbot],
             outputs=chatbot
+        ).then(
+            lambda: "",  # Clear message box
+            outputs=msg
         )
         
         clear_btn.click(