Update app.py

app.py

@@ -2,7 +2,8 @@ import mimetypes
 import os
 import re
 import shutil
-from typing import Optional
+import threading
+from typing import Optional, List, Dict, Any
 
 from dotenv import load_dotenv
 from huggingface_hub import login
@@ -19,7 +20,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,
@@ -58,6 +59,8 @@ 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"
@@ -101,31 +104,52 @@ def setup_environment():
 
 # ------------------------ Model and Tool Management ------------------------
 class ModelManager:
-    """Manages model loading and initialization."""
+    """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]
 
-    @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":
-                return HfApiModel(model_id=model_id)
+                model = HfApiModel(model_id=model_id)
 
             elif chosen_inference == "hf_api_provider":
-                return HfApiModel(provider="together")
+                model = HfApiModel(provider="together")
 
             elif chosen_inference == "litellm":
-                return LiteLLMModel(model_id=model_id)
+                model = LiteLLMModel(model_id=model_id)
 
             elif chosen_inference == "openai":
                 if not key_manager:
                     raise ValueError("Key manager required for OpenAI model")
 
-                return OpenAIServerModel(
+                model = OpenAIServerModel(
                     model_id=model_id, api_key=key_manager.get_key("openai_api_key")
                 )
 
             elif chosen_inference == "transformers":
-                return TransformersModel(
+                model = TransformersModel(
                     model_id="HuggingFaceTB/SmolLM2-1.7B-Instruct",
                     device_map="auto",
                     max_new_tokens=1000,
@@ -134,13 +158,22 @@ 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."""
+    """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)]
 
     @staticmethod
     def load_web_tools(model, browser, text_limit=20000):
@@ -167,51 +200,93 @@ class ToolRegistry:
             )
         except Exception as e:
             print(f"✗ Couldn't initialize image generation tool: {e}")
-            raise
+            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
 
 
 # ------------------------ Agent Creation and Execution ------------------------
-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."
-            )
+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,
+        )
 
-    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,
-    )
+        # Cache the agent for future use
+        self.agent_cache[session_id] = agent
+        return agent
 
 
 def stream_to_gradio(
     agent,
     task: str,
     reset_agent_memory: bool = False,
-    additional_args: Optional[dict] = None,
+    additional_args: Optional[Dict[str, Any]] = None,
 ):
     """Runs an agent with the given task and streams messages as Gradio ChatMessages."""
     for step_log in agent.run(
@@ -247,44 +322,50 @@ def stream_to_gradio(
 
 # ------------------------ Gradio UI Components ------------------------
 class GradioUI:
-    """A one-line interface to launch your agent in Gradio."""
+    """A streamlined interface to launch your agent in Gradio with Zhou Protocol patterns."""
 
-    def __init__(self, file_upload_folder: str | None = None):
+    def __init__(self, file_upload_folder: Optional[str] = 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:
-            if not os.path.exists(file_upload_folder):
-                os.mkdir(file_upload_folder)
+            os.makedirs(self.file_upload_folder, exist_ok=True)
 
     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
-        if "agent" not in session_state:
-            session_state["agent"] = create_agent()
+        agent = self.agent_factory.create_agent(session_state["session_id"])
 
         # Adding monitoring
         try:
             # Log the existence of agent memory
-            has_memory = hasattr(session_state["agent"], "memory")
+            has_memory = hasattr(agent, "memory")
             print(f"Agent has memory: {has_memory}")
             if has_memory:
-                print(f"Memory type: {type(session_state['agent'].memory)}")
+                print(f"Memory type: {type(agent.memory)}")
 
             messages.append(gr.ChatMessage(role="user", content=prompt))
             yield messages
 
-            for msg in stream_to_gradio(
-                session_state["agent"], task=prompt, reset_agent_memory=False
-            ):
+            for msg in stream_to_gradio(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)}")
-            raise
+            messages.append(
+                gr.ChatMessage(
+                    role="assistant", content=f"Error processing request: {str(e)}"
+                )
+            )
+            yield messages
 
     def upload_file(
         self,
@@ -297,6 +378,8 @@ 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
 
@@ -344,7 +427,7 @@ class GradioUI:
         """Process user message and handle file references."""
         message = text_input
 
-        if len(file_uploads_log) > 0:
+        if file_uploads_log and 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 (
@@ -360,7 +443,7 @@ class GradioUI:
     def detect_device(self, request: gr.Request):
         """Detect whether the user is on mobile or desktop device."""
         if not request:
-            return "Unknown device"  # Handle case where request is none.
+            return "Desktop"  # Default to desktop for safety
 
         # Method 1: Check sec-ch-ua-mobile header
         is_mobile_header = request.headers.get("sec-ch-ua-mobile")
@@ -409,7 +492,7 @@ class GradioUI:
             with gr.Sidebar():
                 gr.Markdown(
                     """#OpenDeepResearch - 3theSmolagents!
-                Model_id: google/gemini-2.0-flash-001"""
+                Model_id: R1-1776"""
                 )
                 with gr.Group():
                     gr.Markdown("**What's on your mind mate?**", container=True)
@@ -451,8 +534,9 @@ class GradioUI:
             # Add session state to store session-specific data
             session_state = gr.State({})  # Initialize empty state for each session
             stored_messages = gr.State([])
-            if "file_uploads_log" not in locals():
-                file_uploads_log = gr.State([])
+
+            # Ensure file_uploads_log is always defined
+            file_uploads_log = gr.State([])
 
             chatbot = gr.Chatbot(
                 label="open-Deep-Research",