Update Gradio_UI.py

Gradio_UI.py

@@ -11,102 +11,89 @@ class GradioUI:
         self.search_tool = None
         self.doc_qa_tool = None
 
-        # Try to detect tool types from dictionary-like access
-        tools = getattr(agent, "tools", {})
+        # Ensure agent.tools is iterable
+        tools = agent.tools if isinstance(agent.tools, (list, tuple)) else list(agent.tools)
 
-        for tool in tools.values():  # ✅ Use `.values()` instead of indexing
-            tool_name = getattr(tool, "__name__", None)
-            if tool_name == "get_current_time_in_timezone":
+        for tool in tools:
+            # Match tools based on description (from docstring) or fallback string
+            desc = getattr(tool, "description", "") or str(tool)
+            if "current local time" in desc or "timezone" in desc:
                 self.get_time_tool = tool
-            elif tool_name == "get_current_weather":
+            elif "weather" in desc.lower():
                 self.get_weather_tool = tool
-            elif tool_name == "document_qna_tool":
+            elif "question" in desc.lower() and "pdf" in desc.lower():
                 self.doc_qa_tool = tool
-            elif "text-to-image" in str(tool):
+            elif "text-to-image" in desc.lower() or "image" in str(tool).lower():
                 self.image_gen_tool = tool
-            elif "DuckDuckGoSearchTool" in str(tool):
+            elif "duckduckgo" in str(tool).lower() or "search" in desc.lower():
                 self.search_tool = tool
 
-        # Optional fallback with warnings
+        # 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: search tool"
+        assert self.search_tool, "Missing: DuckDuckGoSearchTool"
 
     def launch(self):
         with gr.Blocks() as demo:
-            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)")
-                    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.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)")
-                    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.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"):
-                    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
-                        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"):
-                    search_query = gr.Textbox(label="Enter search query")
-                    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.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", 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")
-
-                    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.doc_qa_tool(pdf_file.name, question)
-
-                    docqa_btn.click(doc_qa, inputs=[pdf_upload, question_input], outputs=answer_output)
-
-        demo.launch()
+            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()