Gradio_UI.py

import gradio as gr

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 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()