Update app.py

app.py

@@ -2,8 +2,7 @@ import mimetypes
 import os
 import re
 import shutil
-import threading
-from typing import Optional, List, Dict, Any
+from typing import Optional
 
 from dotenv import load_dotenv
 from huggingface_hub import login
@@ -20,7 +19,7 @@ from scripts.text_web_browser import (
     VisitTool,
 )
 from scripts.visual_qa import visualizer
-from scripts.legal_document_tool import LegalDocumentTool
+
 from smolagents import (
     CodeAgent,
     HfApiModel,
@@ -59,8 +58,6 @@ AUTHORIZED_IMPORTS = [
     "fractions",
     "csv",
     "clean-text",
-    "langchain",
-    "llama_index",  # Fixed trailing comma
 ]
 
 user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
@@ -104,52 +101,31 @@ def setup_environment():
 
 # ------------------------ Model and Tool Management ------------------------
 class ModelManager:
-    """Manages model loading and initialization with Zhou Protocol patterns."""
-
-    _instance = None
-    _lock = threading.Lock()
-
-    @classmethod
-    def get_instance(cls):
-        """Thread-safe singleton access to model manager."""
-        if cls._instance is None:
-            with cls._lock:
-                if cls._instance is None:
-                    cls._instance = cls()
-        return cls._instance
-
-    def __init__(self):
-        """Initialize with model cache."""
-        self.model_cache = {}
-
-    def load_model(self, chosen_inference: str, model_id: str, key_manager=None):
-        """Load the specified model with appropriate configuration and caching."""
-        cache_key = f"{chosen_inference}:{model_id}"
-
-        # Return cached model if available
-        if cache_key in self.model_cache:
-            return self.model_cache[cache_key]
+    """Manages model loading and initialization."""
 
+    @staticmethod
+    def load_model(chosen_inference: str, model_id: str, key_manager=None):
+        """Load the specified model with appropriate configuration."""
         try:
             if chosen_inference == "hf_api":
-                model = HfApiModel(model_id=model_id)
+                return HfApiModel(model_id=model_id)
 
             elif chosen_inference == "hf_api_provider":
-                model = HfApiModel(provider="together")
+                return HfApiModel(provider="together")
 
             elif chosen_inference == "litellm":
-                model = LiteLLMModel(model_id=model_id)
+                return LiteLLMModel(model_id=model_id)
 
             elif chosen_inference == "openai":
                 if not key_manager:
                     raise ValueError("Key manager required for OpenAI model")
 
-                model = OpenAIServerModel(
+                return OpenAIServerModel(
                     model_id=model_id, api_key=key_manager.get_key("openai_api_key")
                 )
 
             elif chosen_inference == "transformers":
-                model = TransformersModel(
+                return TransformersModel(
                     model_id="HuggingFaceTB/SmolLM2-1.7B-Instruct",
                     device_map="auto",
                     max_new_tokens=1000,
@@ -158,22 +134,13 @@ class ModelManager:
             else:
                 raise ValueError(f"Invalid inference type: {chosen_inference}")
 
-            # Cache the model for future use
-            self.model_cache[cache_key] = model
-            return model
-
         except Exception as e:
             print(f"✗ Couldn't load model: {e}")
             raise
 
 
 class ToolRegistry:
-    """Manages tool initialization and organization with validation."""
-
-    @staticmethod
-    def validate_tools(tools: List[Tool]) -> List[Tool]:
-        """Validate tools and filter out any None values."""
-        return [tool for tool in tools if isinstance(tool, Tool)]
+    """Manages tool initialization and organization."""
 
     @staticmethod
     def load_web_tools(model, browser, text_limit=20000):
@@ -200,93 +167,51 @@ class ToolRegistry:
             )
         except Exception as e:
             print(f"✗ Couldn't initialize image generation tool: {e}")
-            return None
-
-    @staticmethod
-    def load_legal_document_tool():
-        """Initialize and return the legal document processing tool."""
-        try:
-            # Create a simple instance with default parameters
-            return LegalDocumentTool()
-        except Exception as e:
-            print(f"✗ Couldn't initialize legal document tool: {e}")
-            # Return None instead of raising to make this tool optional
-            return None
+            raise
 
 
 # ------------------------ Agent Creation and Execution ------------------------
-class AgentFactory:
-    """Factory for creating and managing agent instances with Zhou Protocol patterns."""
-
-    _instance = None
-    _lock = threading.Lock()
-
-    @classmethod
-    def get_instance(cls):
-        """Thread-safe singleton access."""
-        if cls._instance is None:
-            with cls._lock:
-                if cls._instance is None:
-                    cls._instance = cls()
-        return cls._instance
-
-    def __init__(self):
-        """Initialize with agent cache."""
-        self.agent_cache = {}
-
-    def create_agent(self, session_id: str = "default") -> CodeAgent:
-        """Creates a fresh agent instance with properly configured tools."""
-        # Return cached agent if available for this session
-        if session_id in self.agent_cache:
-            return self.agent_cache[session_id]
-
-        # Initialize model
-        model = LiteLLMModel(
-            custom_role_conversions=custom_role_conversions,
-            model_id="openrouter/perplexity/r1-1776",  # currently serving:
-        )  # DEEPSEEK = openrouter/perplexity/r1-1776 <--- boss model
-
-        # Initialize tools
-        text_limit = 30000
-        browser = SimpleTextBrowser(**BROWSER_CONFIG)
-
-        # Collect all tools in a single list
-        web_tools = ToolRegistry.load_web_tools(model, browser, text_limit)
-        image_generator = ToolRegistry.load_image_generation_tools()
-        legal_tool = ToolRegistry.load_legal_document_tool()
-
-        # Combine and validate all tools
-        all_tools = [visualizer] + web_tools
-
-        # Only add tools that are properly initialized (not None)
-        if image_generator:
-            all_tools.append(image_generator)
-
-        if legal_tool:
-            all_tools.append(legal_tool)
-
-        # Final validation to ensure all tools are valid
-        all_tools = ToolRegistry.validate_tools(all_tools)
-
-        agent = CodeAgent(
-            model=model,
-            tools=all_tools,  # Pass a single list containing all tools
-            max_steps=10,
-            verbosity_level=1,
-            additional_authorized_imports=AUTHORIZED_IMPORTS,
-            planning_interval=4,
-        )
+def create_agent():
+    """Creates a fresh agent instance with properly configured tools."""
+    # Initialize model
+    model = LiteLLMModel(
+        custom_role_conversions=custom_role_conversions,
+        model_id="openrouter/google/gemini-2.0-flash-001",  # currently serving:
+    )  # DEEPSEEK = openrouter/perplexity/r1-1776 <--- boss model
+
+    # Initialize tools
+    text_limit = 30000
+    browser = SimpleTextBrowser(**BROWSER_CONFIG)
+
+    # Collect all tools in a single list
+    web_tools = ToolRegistry.load_web_tools(model, browser, text_limit)
+    image_generator = ToolRegistry.load_image_generation_tools()
+
+    # Combine all tools into a single list (not a tuple)
+    all_tools = [visualizer] + web_tools + [image_generator]
+
+    # Validate tools before creating agent
+    for tool in all_tools:
+        if not isinstance(tool, Tool):
+            raise ValueError(
+                f"Invalid tool type: {type(tool)}. All tools must be instances of Tool class."
+            )
 
-        # Cache the agent for future use
-        self.agent_cache[session_id] = agent
-        return agent
+    return CodeAgent(
+        model=model,
+        tools=all_tools,  # Pass a single list containing all tools
+        max_steps=10,
+        verbosity_level=1,
+        additional_authorized_imports=AUTHORIZED_IMPORTS,
+        planning_interval=4,
+    )
 
 
 def stream_to_gradio(
     agent,
     task: str,
     reset_agent_memory: bool = False,
-    additional_args: Optional[Dict[str, Any]] = None,
+    additional_args: Optional[dict] = None,
 ):
     """Runs an agent with the given task and streams messages as Gradio ChatMessages."""
     for step_log in agent.run(
@@ -322,50 +247,44 @@ def stream_to_gradio(
 
 # ------------------------ Gradio UI Components ------------------------
 class GradioUI:
-    """A streamlined interface to launch your agent in Gradio with Zhou Protocol patterns."""
+    """A one-line interface to launch your agent in Gradio."""
 
-    def __init__(self, file_upload_folder: Optional[str] = None):
+    def __init__(self, file_upload_folder: str | None = None):
         """Initialize the Gradio UI with optional file upload functionality."""
         self.file_upload_folder = file_upload_folder
-        self.agent_factory = AgentFactory.get_instance()
 
         if self.file_upload_folder is not None:
-            os.makedirs(self.file_upload_folder, exist_ok=True)
+            if not os.path.exists(file_upload_folder):
+                os.mkdir(file_upload_folder)
 
     def interact_with_agent(self, prompt, messages, session_state):
         """Main interaction handler with the agent."""
 
-        # Generate unique session ID if not present
-        if "session_id" not in session_state:
-            session_state["session_id"] = f"session_{id(session_state)}"
-
         # Get or create session-specific agent
-        agent = self.agent_factory.create_agent(session_state["session_id"])
+        if "agent" not in session_state:
+            session_state["agent"] = create_agent()
 
         # Adding monitoring
         try:
             # Log the existence of agent memory
-            has_memory = hasattr(agent, "memory")
+            has_memory = hasattr(session_state["agent"], "memory")
             print(f"Agent has memory: {has_memory}")
             if has_memory:
-                print(f"Memory type: {type(agent.memory)}")
+                print(f"Memory type: {type(session_state['agent'].memory)}")
 
             messages.append(gr.ChatMessage(role="user", content=prompt))
             yield messages
 
-            for msg in stream_to_gradio(agent, task=prompt, reset_agent_memory=False):
+            for msg in stream_to_gradio(
+                session_state["agent"], task=prompt, reset_agent_memory=False
+            ):
                 messages.append(msg)
                 yield messages  # Yield messages after each step
             yield messages  # Yield messages one last time
 
         except Exception as e:
             print(f"Error in interaction: {str(e)}")
-            messages.append(
-                gr.ChatMessage(
-                    role="assistant", content=f"Error processing request: {str(e)}"
-                )
-            )
-            yield messages
+            raise
 
     def upload_file(
         self,
@@ -378,8 +297,6 @@ class GradioUI:
 
         try:
             mime_type, _ = mimetypes.guess_type(file.name)
-            if not mime_type:
-                return gr.Textbox("Unknown file type", visible=True), file_uploads_log
         except Exception as e:
             return gr.Textbox(f"Error: {e}", visible=True), file_uploads_log
 
@@ -427,7 +344,7 @@ class GradioUI:
         """Process user message and handle file references."""
         message = text_input
 
-        if file_uploads_log and len(file_uploads_log) > 0:
+        if len(file_uploads_log) > 0:
             message += f"\nYou have been provided with these files, which might be helpful or not: {file_uploads_log}"  # Added file list
 
         return (
@@ -443,7 +360,7 @@ class GradioUI:
     def detect_device(self, request: gr.Request):
         """Detect whether the user is on mobile or desktop device."""
         if not request:
-            return "Desktop"  # Default to desktop for safety
+            return "Unknown device"  # Handle case where request is none.
 
         # Method 1: Check sec-ch-ua-mobile header
         is_mobile_header = request.headers.get("sec-ch-ua-mobile")
@@ -492,7 +409,7 @@ class GradioUI:
             with gr.Sidebar():
                 gr.Markdown(
                     """#OpenDeepResearch - 3theSmolagents!
-                Model_id: R1-1776"""
+                Model_id: google/gemini-2.0-flash-001"""
                 )
                 with gr.Group():
                     gr.Markdown("**What's on your mind mate?**", container=True)
@@ -534,9 +451,8 @@ class GradioUI:
             # Add session state to store session-specific data
             session_state = gr.State({})  # Initialize empty state for each session
             stored_messages = gr.State([])
-
-            # Ensure file_uploads_log is always defined
-            file_uploads_log = gr.State([])
+            if "file_uploads_log" not in locals():
+                file_uploads_log = gr.State([])
 
             chatbot = gr.Chatbot(
                 label="open-Deep-Research",