Update Gradio_UI.py

Gradio_UI.py

@@ -4,99 +4,82 @@ 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
-
-        # Ensure agent.tools is iterable
-        tools = agent.tools if isinstance(agent.tools, (list, tuple)) else list(agent.tools)
-
-        for tool in tools:
-            # Get raw function name if possible
-            raw_func = getattr(tool, "fn", None)
-            raw_func_name = getattr(raw_func, "__name__", "").lower()
-            tool_str = str(tool).lower()
-        
-            if "timezone" in raw_func_name or "get_current_time_in_timezone" in raw_func_name:
-                self.get_time_tool = tool
-            elif "weather" in raw_func_name or "get_current_weather" in raw_func_name:
-                self.get_weather_tool = tool
-            elif "document_qna_tool" in raw_func_name:
-                self.doc_qa_tool = tool
-            elif "text-to-image" in tool_str or "image" in raw_func_name:
-                self.image_gen_tool = tool
-            elif "duckduckgo" in tool_str or "search" in raw_func_name:
-                self.search_tool = tool
-
-        # Validate all tools are loaded
-        assert self.get_time_tool, "Missing: get_current_time_in_timezone"
-        assert self.get_weather_tool, "Missing: get_current_weather"
-        assert self.doc_qa_tool, "Missing: document_qna_tool"
-        assert self.image_gen_tool, "Missing: image generation tool"
-        assert self.search_tool, "Missing: DuckDuckGoSearchTool"
-
     def launch(self):
         with gr.Blocks() as demo:
-            gr.Markdown("# 🧠 Document Q&A + AI Agent Interface")
-
-            with gr.Tab("📄 Document Q&A"):
-                pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
-                question_input = gr.Textbox(label="Ask a question about the PDF")
-                docqa_output = gr.Textbox(label="Answer")
-                docqa_btn = gr.Button("Get Answer")
-
-            with gr.Tab("🌦️ Weather"):
-                place_input = gr.Textbox(label="Enter a city")
-                weather_output = gr.Textbox(label="Weather Info")
-                weather_btn = gr.Button("Get Weather")
-
-            with gr.Tab("🕓 Time"):
-                timezone_input = gr.Textbox(label="Enter timezone (e.g., 'Asia/Kolkata')")
-                time_output = gr.Textbox(label="Current Time")
-                time_btn = gr.Button("Get Time")
-
-            with gr.Tab("🖼️ Image Generator"):
-                prompt_input = gr.Textbox(label="Enter image prompt")
-                image_output = gr.Image(label="Generated Image")
-                image_btn = gr.Button("Generate Image")
-
-            with gr.Tab("🔍 Web Search"):
-                search_input = gr.Textbox(label="Enter search query")
-                search_output = gr.Textbox(label="Search Result")
-                search_btn = gr.Button("Search")
-
-            # Actions
-            docqa_btn.click(
-                fn=lambda pdf, q: self.doc_qa_tool(pdf.name, q),
-                inputs=[pdf_input, question_input],
-                outputs=docqa_output
-            )
-
-            weather_btn.click(
-                fn=lambda place: self.get_weather_tool(place),
-                inputs=place_input,
-                outputs=weather_output
-            )
-
-            time_btn.click(
-                fn=lambda tz: self.get_time_tool(tz),
-                inputs=timezone_input,
-                outputs=time_output
-            )
-
-            image_btn.click(
-                fn=lambda prompt: self.image_gen_tool(prompt),
-                inputs=prompt_input,
-                outputs=image_output
-            )
-
-            search_btn.click(
-                fn=lambda query: self.search_tool(query),
-                inputs=search_input,
-                outputs=search_output
-            )
-
-        demo.launch()
+            gr.Markdown("# Multi-Tool AI Agent")
+            gr.Markdown("Use the tabs below to interact with 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)
+                    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)
+
+                    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)
+                    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)
+
+                    get_time_btn.click(get_time, inputs=timezone_input, outputs=time_output)
+
+                # Tab 3: 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")
+
+                    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
+
+                    gen_image_btn.click(gen_image, inputs=image_prompt, outputs=image_output)
+
+                # Tab 4: 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_btn = gr.Button("Search")
+
+                    def search(q):
+                        if not q.strip():
+                            return "Please enter a search query."
+                        return self.agent.tools[3](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")
+                    answer_output = gr.Textbox(label="Answer", interactive=False)
+                    docqa_btn = gr.Button("Get Answer")
+
+                    def doc_qa(pdf_file, question):
+                        if pdf_file is None:
+                            return "Please upload a PDF file."
+                        if not question.strip():
+                            return "Please enter a question."
+                        return self.agent.tools[4](pdf_file.name, question)
+
+                    docqa_btn.click(doc_qa, inputs=[pdf_upload, question_input], outputs=answer_output)
+
+            demo.launch()