Update Gradio_UI.py

Gradio_UI.py

@@ -4,41 +4,72 @@ class GradioUI:
     def __init__(self, agent):
         self.agent = agent
 
+        # Initialize tool references
+        self.get_time_tool = None
+        self.get_weather_tool = None
+        self.image_gen_tool = None
+        self.search_tool = None
+        self.doc_qa_tool = None
+
+        # Map tools by known types or names
+        for tool in self.agent.tools:
+            try:
+                tool_name = getattr(tool, "__name__", None)
+                if tool_name == "get_current_time_in_timezone":
+                    self.get_time_tool = tool
+                elif tool_name == "get_current_weather":
+                    self.get_weather_tool = tool
+                elif tool_name == "document_qna_tool":
+                    self.doc_qa_tool = tool
+                elif "text-to-image" in str(tool):
+                    self.image_gen_tool = tool
+                elif "DuckDuckGoSearchTool" in str(tool):
+                    self.search_tool = tool
+            except Exception:
+                pass
+
+        # Fallback in case mapping fails
+        self.get_time_tool = self.get_time_tool or self.agent.tools[0]
+        self.get_weather_tool = self.get_weather_tool or self.agent.tools[1]
+        self.image_gen_tool = self.image_gen_tool or self.agent.tools[2]
+        self.search_tool = self.search_tool or self.agent.tools[3]
+        self.doc_qa_tool = self.doc_qa_tool or self.agent.tools[4]
+
     def launch(self):
         with gr.Blocks() as demo:
-            gr.Markdown("# Multi-Tool AI Agent")
-            gr.Markdown("Use the tabs below to interact with different tools:")
+            gr.Markdown("# 🤖 Multi-Tool AI Agent")
+            gr.Markdown("Use the tabs below to interact with the different tools:")
 
             with gr.Tabs():
 
                 # Tab 1: Weather
-                with gr.TabItem("Weather"):
-                    place_input = gr.Textbox(label="Enter place name (e.g., London)", placeholder="City or location")
-                    weather_output = gr.Textbox(label="Current Weather", interactive=False)
+                with gr.TabItem("🌦️ Weather"):
+                    place_input = gr.Textbox(label="Enter place name (e.g., London)")
+                    weather_output = gr.Textbox(label="Weather Info", interactive=False)
                     get_weather_btn = gr.Button("Get Weather")
 
                     def get_weather(place):
                         if not place.strip():
                             return "Please enter a valid place name."
-                        return self.agent.tools[1](place)
+                        return self.get_weather_tool(place)
 
                     get_weather_btn.click(get_weather, inputs=place_input, outputs=weather_output)
 
                 # Tab 2: Local Time
-                with gr.TabItem("Local Time"):
-                    timezone_input = gr.Textbox(label="Enter timezone (e.g., America/New_York)", placeholder="Timezone string")
-                    time_output = gr.Textbox(label="Current Time", interactive=False)
+                with gr.TabItem("🕒 Local Time"):
+                    timezone_input = gr.Textbox(label="Enter timezone (e.g., America/New_York)")
+                    time_output = gr.Textbox(label="Local Time", interactive=False)
                     get_time_btn = gr.Button("Get Local Time")
 
                     def get_time(tz):
                         if not tz.strip():
                             return "Please enter a valid timezone."
-                        return self.agent.tools[0](tz)
+                        return self.get_time_tool(tz)
 
                     get_time_btn.click(get_time, inputs=timezone_input, outputs=time_output)
 
                 # Tab 3: Image Generation
-                with gr.TabItem("Image Generation"):
+                with gr.TabItem("🎨 Image Generation"):
                     image_prompt = gr.Textbox(label="Enter image description prompt", lines=2)
                     image_output = gr.Image(label="Generated Image")
                     gen_image_btn = gr.Button("Generate Image")
@@ -46,30 +77,27 @@ class GradioUI:
                     def gen_image(prompt):
                         if not prompt.strip():
                             return None
-                        # The image generation tool might return an image URL or PIL.Image.
-                        # Adjust this depending on your tool's output format.
-                        result = self.agent.tools[2](prompt)
-                        return result
+                        return self.image_gen_tool(prompt)
 
                     gen_image_btn.click(gen_image, inputs=image_prompt, outputs=image_output)
 
                 # Tab 4: Web Search
-                with gr.TabItem("Web Search"):
+                with gr.TabItem("🔍 Web Search"):
                     search_query = gr.Textbox(label="Enter search query")
-                    search_output = gr.Textbox(label="Search Results", interactive=False)
+                    search_output = gr.Textbox(label="Search Result", interactive=False)
                     search_btn = gr.Button("Search")
 
                     def search(q):
                         if not q.strip():
                             return "Please enter a search query."
-                        return self.agent.tools[3](q)
+                        return self.search_tool(q)
 
                     search_btn.click(search, inputs=search_query, outputs=search_output)
 
                 # Tab 5: Document Q&A
-                with gr.TabItem("Document Q&A"):
-                    pdf_upload = gr.File(label="Upload PDF Document", file_types=[".pdf"])
-                    question_input = gr.Textbox(label="Enter your question about the document")
+                with gr.TabItem("📄 Document Q&A"):
+                    pdf_upload = gr.File(label="Upload PDF", file_types=[".pdf"])
+                    question_input = gr.Textbox(label="Ask a question about the document")
                     answer_output = gr.Textbox(label="Answer", interactive=False)
                     docqa_btn = gr.Button("Get Answer")
 
@@ -78,8 +106,8 @@ class GradioUI:
                             return "Please upload a PDF file."
                         if not question.strip():
                             return "Please enter a question."
-                        return self.agent.tools[4](pdf_file.name, question)
+                        return self.doc_qa_tool(pdf_file.name, question)
 
                     docqa_btn.click(doc_qa, inputs=[pdf_upload, question_input], outputs=answer_output)
 
-            demo.launch()
+        demo.launch()