System: Generate search queries based on the following query and context.
            use the following format for search queries:

```search
{
  "queries": [
    "first search query",
    "second search query",
    "third search query",
  ]
}```
            
User: Query: Make an improved version of this code, only write FULL CODE in one python file : import os
import asyncio
import gradio as gr
from collections import defaultdict
import time
from concurrent.futures import ThreadPoolExecutor
import json
import re
from typing import List, Dict, Any
from langdetect import detect
from translate import Translator
import matplotlib.pyplot as plt
import io
import base64
from PIL import Image, ImageDraw
import aiohttp
from bs4 import BeautifulSoup
from duckduckgo_search import DDGS
from gradio_client import Client

# Constants
DUCKDUCKGO_BASE_URL = "https://html.duckduckgo.com/html/?q="

# Initialize ThreadPoolExecutor
executor = ThreadPoolExecutor(max_workers=3)

# Enhanced Prompts
class Prompts:
    initial = """You are an **advanced AI assistant** capable of detailed, multi-step problem solving. Break down tasks into clear, logical steps. Explain each step thoroughly and provide a final answer. If you need additional information, use the following format for search queries IMPORTANT:

```search
{
  "queries": [
    "first search query",
    "second search query",
    "third search query",
    "etc"
  ]
}```
"""
    followup = """You are an **AI refiner** tasked with improving previous solutions. Review past attempts, provide new insights, and enhance the solution step by step.  If you need additional information, use the following format for search queries IMPORTANT:

```search
{
  "queries": [
    "first search query",
    "second search query",
    "third search query",
    "etc"
  ]
}```
"""
    contrastive = """You are an **AI challenger** tasked with providing a contrasting perspective to previous solutions. Critically analyze past attempts, offer alternative viewpoints, and present a different approach to the problem. Challenge assumptions and explore unconventional solutions.   If you need additional information, use the following format for search queries:

```search
{
  "queries": [
    "first search query",
    "second search query",
    "third search query",
    "etc"
  ]
}```
"""
    synthesis = "As an **AI synthesizer**, combine multiple solution attempts into a comprehensive answer. Compare methods, synthesize key insights, and provide a final, refined and complete answer."
    synthetic_knowledge = "Generate **synthetic information** based on existing data to aid in problem solving. Use synthetic examples and clearly label them as such."

# Enhanced Knowledge Base
class KnowledgeBase:
   def __init__(self):
       self.base = defaultdict(list)

   def add_information(self, topic: str, data: str) -> None:
       self.base[topic.lower()].append(data)

   def fetch_information(self, topic: str) -> list:
       return self.base.get(topic.lower(), [])

   def export_to_json(self, filename: str) -> None:
       with open(filename, 'w') as f:
           json.dump(dict(self.base), f)

   def import_from_json(self, filename: str) -> None:
       with open(filename, 'r') as f:
           data = json.load(f)
           self.base = defaultdict(list, data)

# Enhanced Web API Client
class WebApiClient:
   def __init__(self):
       self.base_url = DUCKDUCKGO_BASE_URL
       self.headers = {
           "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
       }

   async def fetch_web_information(self, query: str) -> list:
       try:
           async with aiohttp.ClientSession() as session:
               async with session.post(self.base_url, data={'q': query}, headers=self.headers) as response:
                   if response.status == 200:
                       html_content = await response.text()
                       soup = BeautifulSoup(html_content, 'html.parser')
                       result_links = soup.find_all('a', class_='result__a')
                       snippets = [link.get_text() for link in result_links]
                       return snippets[:5]
                   else:
                       print(f"DuckDuckGo returned status code {response.status}")
                       return []
       except Exception as e:
           print(f"Error fetching web information from DuckDuckGo: {e}")
           return []

   async def fetch_url_content(self, url: str) -> str:
       try:
           async with aiohttp.ClientSession() as session:
               async with session.get(url, headers=self.headers, timeout=10) as response:
                   if response.status == 200:
                       html_content = await response.text()
                       soup = BeautifulSoup(html_content, 'html.parser')
                       text_content = ' '.join([tag.get_text() for tag in soup.find_all(['p', 'h1', 'h2', 'h3', 'li'])])
                       return text_content[:1000]
                   else:
                       return f"Error: Unable to fetch content from {url}"
       except Exception as e:
           return f"Error fetching content from {url}: {e}"

# Enhanced Model Client using gradio_client
class ModelClient:
    def __init__(self):
        self.client = Client("yuntian-deng/o1mini")
        self.web_client = WebApiClient()
    
    async def call_model(self, messages: list, temp: float = 0.1, tokens: int = 8000) -> str:
        # Format the messages into a prompt string
        prompt = ""
        for message in messages:
            role = message.get("role", "")
            content = message.get("content", "")
            if role == "system":
                prompt += f"System: {content}\n"
            elif role == "user":
                prompt += f"User: {content}\n"
            elif role == "assistant":
                prompt += f"Assistant: {content}\n"
        prompt += "Assistant:"

        # Prepare the inputs
        inputs = prompt

        # Call the model
        try:
            # Since gradio_client's predict is synchronous, use asyncio.to_thread to run it asynchronously
            result = await asyncio.to_thread(
                self.client.predict,
                inputs=inputs,
                top_p=1,
                temperature=temp,
                chat_counter=0,
                chatbot=[],
                api_name="/predict"
            )

            
            clean_result = "".join(result[0][-1][-1]) #.strip()
            clean_result = post_process_response(clean_result)
            print(clean_result)
            
            return clean_result
        except Exception as e:
            print(f"Error calling o1mini model: {e}")
            return "I'm sorry, but I encountered an error while processing your request."

    async def generate_synthetic_information(self, query: str, related_info: list) -> str:
        messages = [
            {"role": "system", "content": Prompts.synthetic_knowledge},
            {"role": "user", "content": f"Query: {query}\n\nExisting Info:\n{chr(10).join(related_info) if related_info else 'None'}"}
        ]
        return await self.call_model(messages)

    async def retrieve_url_content(self, url: str) -> str:
        return await self.web_client.fetch_url_content(url)

    async def process_with_url_retrieval(self, response: str, context: list = None) -> str:
        if "URL_RETRIEVAL_NEEDED:" in response:
            url = response.split("URL_RETRIEVAL_NEEDED:", 1)[1].strip()
            url_content = await self.retrieve_url_content(url)
            
            messages = [
                {"role": "system", "content": "You are an AI assistant tasked with providing a refined answer based on the original response and newly retrieved information."},
                {"role": "user", "content": f"Original response:\n\n{response}\n\nRetrieved content from {url}:\n\n{url_content}\n\nPlease provide a refined and more accurate answer based on this additional information."}
            ]
            
            refined_response = await self.call_model(messages)
            return f"Refined answer based on data from {url}:\n\n{refined_response}"
        else:
            return response

def post_process_response(text: str) -> str:
    """
    Cleans the LLM response by removing Markdown headers and HTML tags.

    Args:
        text (str): The raw response from the LLM.

    Returns:
        str: The cleaned text.
    """
    # Remove Markdown headers (e.g., ### Heading)
    text = re.sub(r'^#{1,6}\s*', '', text, flags=re.MULTILINE)

    # Remove horizontal rules
    text = re.sub(r'^---$', '', text, flags=re.MULTILINE)

    # Remove HTML tags
    soup = BeautifulSoup(text, 'html.parser')
    text = soup.get_text(separator='\n')  # You can choose a different separator if needed

    # Optionally, remove extra newlines and whitespace
    text = re.sub(r'\n{2,}', '\n\n', text)
    #text = text.strip()

    return text

# Enhanced LLM Search Extension
class LLMSearchExtension:
   def __init__(self):
       self.search_engine = SearchEngineManager()
       self.result_processor = ResultProcessor()
       self.content_retriever = ContentRetriever()
       self.output_formatter = OutputFormatter()

   async def handle_query(self, query, context=None, options=None):
       options = options or {}
       raw_results = await self.search_engine.execute_search(query, options)
       processed_results = self.result_processor.process(raw_results, context)
       enriched_results = await self.content_retriever.retrieve(processed_results, options)
       formatted_results = self.output_formatter.format(enriched_results)
       return formatted_results

# Enhanced Search Layer
class DuckDuckGoSearch:
   async def search(self, query, options):
       max_results = options.get('max_results', 5)
       loop = asyncio.get_event_loop()
       results = await loop.run_in_executor(
           executor,
           lambda: list(DDGS().text(query, max_results=max_results))
       )
       return results

class SearchEngineManager:
   def __init__(self):
       self.engines = {
           'duckduckgo': DuckDuckGoSearch(),
       }

   async def execute_search(self, query, options):
       engine_name = options.get('engine', 'duckduckgo')
       search_engine = self.engines.get(engine_name)
       if not search_engine:
           raise ValueError(f"Search engine '{engine_name}' not supported.")
       return await search_engine.search(query[:50], options)

# Enhanced Processing Layer
class ResultProcessor:
   def process(self, raw_results, context):
       processed_results = []
       for result in raw_results:
           summary = self.summarize(result.get('body', '') if isinstance(result, dict) else result, context)
           processed_results.append({
               'title': result.get('title', '') if isinstance(result, dict) else result,
               'summary': summary,
               'url': result.get('href', '') if isinstance(result, dict) else '',
           })
       return processed_results

   def summarize(self, text, context):
       return text[:200] + '...'

# Enhanced Content Retrieval Layer
class ContentRetriever:
   async def retrieve(self, processed_results, options):
       relevant_tags = options.get('relevant_tags', ['p', 'h1', 'h2'])
       tasks = [self.process_result(result, relevant_tags) for result in processed_results]
       enriched_results = await asyncio.gather(*tasks)
       return enriched_results

   async def process_result(self, result, tags):
       url = result['url']
       if not url:
           result['extracted_content'] = "No URL provided."
           return result
       try:
           page_content = await self.fetch_page(url)
           extracted_content = self.extract_content(page_content, tags)
           result['extracted_content'] = extracted_content
       except Exception as e:
           result['extracted_content'] = f"Error fetching content: {e}"
       return result

   async def fetch_page(self, url):
       headers = {'User-Agent': 'Mozilla/5.0'}
       async with aiohttp.ClientSession() as session:
           async with session.get(url, headers=headers, timeout=10) as response:
               response.raise_for_status()
               return await response.text()

   def extract_content(self, html_content, tags):
       soup = BeautifulSoup(html_content, 'html.parser')
       extracted_texts = []
       for tag in tags:
           for element in soup.find_all(tag):
               text = element.get_text(strip=True)
               if text:
                   extracted_texts.append(text)
       return ' '.join(extracted_texts)

# Enhanced Output Layer
class OutputFormatter:
   def format(self, enriched_results):
       return enriched_results

# Enhanced Processing Steps
class ProcessingSteps:
   def __init__(self, model_client: ModelClient, knowledge_base: KnowledgeBase, search_extension: LLMSearchExtension):
       self.model_client = model_client
       self.knowledge_base = knowledge_base
       self.search_extension = search_extension
       self.translator = Translator(to_lang="en")

   async def generate_search_tags(self, query: str, context: list = None) -> list:
       messages = [
           {"role": "system", "content": """Generate search queries based on the following query and context.
            use the following format for search queries:

```search
{
  "queries": [
    "first search query",
    "second search query",
    "third search query",
  ]
}```
            """},
           {"role": "user", "content": f"Query: {query}\nContext: {chr(10).join(context) if context else 'None'}"}
       ]
       response = await self.model_client.call_model(messages)
       response = json.loads(response)
       tags = [tag.strip() for tag in response["queries"] if tag.strip()]
       return tags

   async def processing_step(self, query: str, prev_steps: list = None, current_iteration: int = 1, progress: gr.Progress = None) -> (list, list):
        is_first = current_iteration == 1
        relevant_info = []
        all_tags = set()

        context = prev_steps if not is_first else []

        # Detect language and translate if necessary
        detected_lang = detect(query)
        if detected_lang != 'en':
            query = self.translator.translate(query)

        # URL retrieval step
        url_retrieval_response = await self.model_client.process_with_url_retrieval(query, context)
        if "Refined answer based on data from " in url_retrieval_response:
            relevant_info.append(url_retrieval_response)

        search_tags = await self.generate_search_tags(query, context)
        all_tags.update(search_tags)

        search_tags = search_tags[:3]
        for tag in search_tags:
            search_options = {
                'engine': 'duckduckgo',
                'max_results': 3,
                'relevant_tags': ['p', 'h1', 'h2']
            }
            search_results = await self.search_extension.handle_query(tag, context=None, options=search_options)
            for res in search_results:
                if res.get('summary'):
                    relevant_info.append(res['summary'])

        print(":::relevantinfofromweb:::")
        print(relevant_info)

        for topic in self.knowledge_base.base:
            if topic in query.lower():
                relevant_info.extend(self.knowledge_base.fetch_information(topic))

        if not relevant_info and not is_first:
            synthetic_info = await self.model_client.generate_synthetic_information(query, related_info=[])
            relevant_info.append(f"Synthetic Information: {synthetic_info}")

        if is_first:
            system_prompt = Prompts.initial
        elif current_iteration in [2, 4]:
            system_prompt = Prompts.contrastive
        else:
            system_prompt = Prompts.followup

        prev_content = "\n\n".join(prev_steps) if prev_steps else ""
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Query: {query}\n\nIteration: {current_iteration}\n\nPrevious Attempts:\n{prev_content}\n\nRelevant Info:\n{chr(10).join(relevant_info) if relevant_info else 'None'}"}
        ]

        step_response = await self.model_client.call_model(messages)

        # Check for immediate URL retrieval
        if "URL_RETRIEVAL_NEEDED:" in step_response:
            step_response = await self.model_client.process_with_url_retrieval(step_response, context)

        steps = prev_steps or []
        steps.append(step_response)

        new_search_queries = self.extract_search_commands(step_response)
        for search_query in new_search_queries:
            search_results = await self.search_extension.handle_query(search_query, context=None, options={'engine': 'duckduckgo', 'max_results': 3, 'relevant_tags': ['p', 'h1', 'h2']})
            for res in search_results:
                if res.get('summary'):
                    relevant_info.append(f"Search for '{search_query}': {res['summary']}")

        new_tags = await self.generate_search_tags(query, steps)
        all_tags.update(new_tags)

        return list(all_tags), steps

   
   def extract_search_commands(self, response: str) -> list:
        search_commands = []
        
        # Use the updated regex to find JSON-like blocks in the response
        search_blocks = re.findall(r'{\s*"queries":\s*\[.*?\]\s*}', response, re.DOTALL)
        
        # Loop through each block and try to parse it as JSON
        for block in search_blocks:
            try:
                # Load the block as JSON
                search_data = json.loads(block)
                
                # Ensure it's a dictionary and contains the 'queries' key
                if isinstance(search_data, dict) and 'queries' in search_data:
                    # Extend the search_commands list with the queries found in the block
                    search_commands.extend(search_data['queries'])
            except json.JSONDecodeError:
                # Handle the case where the block isn't valid JSON
                print(f"Error parsing search block: {block}")
        
        # Return the list of search commands (queries)
        return search_commands

   async def synthesize_steps(self, query: str, all_steps: list) -> str:
       steps_text = "\n\n".join([f"Attempt {i+1}:\n{step}" for i, step in enumerate(all_steps)])
       messages = [
           {"role": "system", "content": Prompts.synthesis},
           {"role": "user", "content": f"Original Query: {query}\n\nAll Attempts:\n{steps_text}"}
       ]
       return await self.model_client.call_model(messages)

# Modify the full_processing function
async def full_processing(query: str, progress: gr.Progress) -> tuple:
    start_time = time.time()
    steps = []
    final_output = f"# Detailed Processing\n\n## Query\n{query}\n\n"

    max_iterations = 5
    current_iteration = 0
    new_tags = None

    metrics = {
        'Iterations': 0,
        'Total Attempt Time (s)': 0.0,
        'Synthesis Time (s)': 0.0,
        'Total Time (s)': 0.0,
        'Web Searches': 0,
        'Knowledge Base Accesses': 0
    }

    while current_iteration < max_iterations:
        current_iteration += 1
        metrics['Iterations'] = current_iteration
        progress_step = (current_iteration) / (max_iterations + 1)
        progress(progress_step, desc=f"Processing iteration {current_iteration}...")

        step_start = time.time()
        new_tags, steps = await processing_steps.processing_step(query, steps, current_iteration, progress=progress)
        step_duration = time.time() - step_start
        metrics['Total Attempt Time (s)'] += step_duration

        metrics['Web Searches'] += len(processing_steps.extract_search_commands(steps[-1]))
        metrics['Knowledge Base Accesses'] += sum(1 for topic in knowledge_base.base if topic in query.lower())

        approach_type = "Contrastive" if current_iteration in [2, 4] else "Regular"
        final_output += f"## Attempt {current_iteration} ({approach_type} Approach)\n{steps[-1]}\n\n"

        # Check for URL retrieval after generation
        if "Refined answer based on data from " in steps[-1]:
            progress(progress_step, desc="Retrieving additional information...")
            refined_response = await processing_steps.model_client.process_with_url_retrieval(steps[-1], context=steps)
            steps[-1] = refined_response
            final_output = final_output.rsplit("##", 1)[0]  # Remove the last attempt
            final_output += f"## Attempt {current_iteration} (Refined with URL data)\n{refined_response}\n\n"

        # Update UI after each attempt
        yield update_ui(final_output, metrics)

        if not new_tags and current_iteration >= 4:
            break

    progress(0.9, desc="Synthesizing final answer...")
    synthesis_start = time.time()
    synthesis = await processing_steps.synthesize_steps(query, steps)
    synthesis_duration = time.time() - synthesis_start
    metrics['Synthesis Time (s)'] += synthesis_duration
    final_output += f"## Final Synthesis\n{synthesis}\n\n"

    # Check for URL retrieval after final synthesis
    if "Refined answer based on data from " in synthesis:
        progress(0.95, desc="Retrieving additional information for final refinement...")
        refined_synthesis = await processing_steps.model_client.process_with_url_retrieval(synthesis, context=steps)
        synthesis = refined_synthesis
        final_output = final_output.rsplit("##", 1)[0]  # Remove the last synthesis
        final_output += f"## Final Synthesis (Refined with URL data)\n{refined_synthesis}\n\n"
    
    end_time = time.time()
    metrics['Total Time (s)'] = end_time - start_time

    final_output += generate_timing_info(metrics)
    
    yield update_ui(final_output, metrics)

def generate_timing_info(metrics):
   return f"""## Timing Information
- Number of Iterations: {metrics['Iterations']}
- Time for Attempts: {metrics['Total Attempt Time (s)']:.2f}s
- Synthesis Time: {metrics['Synthesis Time (s)']:.2f}s
- Total Time: {metrics['Total Time (s)']:.2f}s
- Web Searches: {metrics['Web Searches']}
- Knowledge Base Accesses: {metrics['Knowledge Base Accesses']}
"""

def update_ui(final_output, metrics):
   metrics_markdown = f"""
   ## Processing Metrics
   - **Iterations**: {metrics['Iterations']}
   - **Time for Attempts**: {metrics['Total Attempt Time (s)']:.2f}s
   - **Synthesis Time**: {metrics['Synthesis Time (s)']:.2f}s
   - **Total Time**: {metrics['Total Time (s)']:.2f}s
   - **Web Searches**: {metrics['Web Searches']}
   - **Knowledge Base Accesses**: {metrics['Knowledge Base Accesses']}
   """

   graph_base64 = generate_performance_graph(metrics)
   
   return final_output, metrics_markdown, f"data:image/png;base64,{graph_base64}"

def generate_performance_graph(metrics):
   fig, ax = plt.subplots()
   ax.bar(metrics.keys(), metrics.values())
   ax.set_title('Processing Performance Metrics')
   ax.set_ylabel('Value')
   plt.xticks(rotation=45, ha='right')
   plt.tight_layout()
   
   buf = io.BytesIO()
   plt.savefig(buf, format='png')
   buf.seek(0)
   
   return base64.b64encode(buf.getvalue()).decode()

def capture_screenshot(html_content):
   # Convert HTML to image (this is a simplified version, you might need to use a more robust solution)
   img = Image.new('RGB', (800, 600), color = (255, 255, 255))
   d = ImageDraw.Draw(img)
   d.text((10,10), html_content, fill=(0,0,0))
   
   buffered = io.BytesIO()
   img.save(buffered, format="PNG")
   img_str = base64.b64encode(buffered.getvalue()).decode()
   return f"data:image/png;base64,{img_str}"

# Modify the gradio_handler function
async def gradio_handler(query: str, progress: gr.Progress = gr.Progress()):
    output_generator = full_processing(query, progress)
    final_output = ""
    first_generation = True
    async for output, metrics_markdown, graph_data_url in output_generator:
        final_output = output
        loader_visibility = gr.update(visible=False) if first_generation else gr.update()
        yield output, metrics_markdown, gr.update(value=graph_data_url, visible=True), loader_visibility
        first_generation = False
    
    # Final yield to ensure the last update is sent
    yield final_output, metrics_markdown, gr.update(value=graph_data_url, visible=True), gr.update(visible=False)

# Enhanced Custom CSS for styling
custom_css = """
.gradio-container {
   font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
   background-color: #f5f7fa;
   color: #333;
   padding: 20px;
}

h1 {
   color: #4a90e2;
   text-align: center;
   margin-bottom: 20px;
}

#query-input textarea {
   border: 2px solid #4a90e2;
   border-radius: 8px;
   padding: 10px;
   font-size: 16px;
   resize: vertical;
}

#submit-button, #screenshot-button {
   background-color: #4a90e2;
   color: white;
   font-size: 18px;
   padding: 10px 20px;
   border-radius: 8px;
   border: none;
   cursor: pointer;
   transition: background-color 0.3s ease;
}

#submit-button:hover, #screenshot-button:hover {
   background-color: #357abd;
}

#output-display, #metrics-display {
   background-color: #ffffff;
   border: 2px solid #4a90e2;
   border-radius: 8px;
   padding: 20px;
   font-size: 16px;
   overflow: auto !important;
   max-height: 600px;
   margin-top: 20px;
   box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}


#loader-overlay {
    display:none;
    position: fixed;
    top: 0;
    left: 0;
    width: 0%;
    height: 0%;
    background-color: rgba(255, 255, 255, 0.8);
    display: flex;
    justify-content: center;
    align-items: center;
    z-index: 1000;
}
.pending{
opacity:1 !important;
}
.loader {
    border: 8px solid #f3f3f3;
    border-top: 8px solid #4a90e2;
    border-radius: 50%;
    width: 60px;
    height: 60px;
    animation: spin 1s linear infinite;
}

@keyframes spin {
   0% { transform: rotate(0deg); }
   100% { transform: rotate(360deg); }
}

@media only screen and (max-width: 768px) {
   .gradio-container .row {
       flex-direction: column;
   }

   #output-display, #metrics-display {
       max-height: 400px;
   }
}
"""

# Initialize components
knowledge_base = KnowledgeBase()
knowledge_base.add_information("python", "Python is a versatile programming language used for various applications.")
knowledge_base.add_information("asyncio", "Asyncio is a library to write concurrent code using the async/await syntax.")

search_extension = LLMSearchExtension()

# Initialize ModelClient without GROQ_API_KEY
model_client = ModelClient()

processing_steps = ProcessingSteps(model_client, knowledge_base, search_extension)

# Create the Enhanced Gradio interface
with gr.Blocks(css=custom_css) as interface:
    gr.Markdown("""
    # 🧠 Advanced Multi-Step Processing Assistant
    **Enter a query to see detailed multi-step processing with integrated knowledge, web information retrieval, dynamic search capabilities, and synthesis.**
    """)

    with gr.Row():
        with gr.Column(scale=1):
            query_input = gr.Textbox(
                label="Enter your query",
                placeholder="e.g., What is the tallest mountain in the world?",
                lines=2,
                max_lines=5,
                elem_id="query-input"
            )
            submit_button = gr.Button("Submit", elem_id="submit-button")
        with gr.Column(scale=1):
            metrics_display = gr.Markdown(
                label="## Processing Metrics",
                value="Awaiting input...",
                elem_id="metrics-display"
            )

    with gr.Row():
        output_display = gr.Markdown(
            label="## Processing Steps",
            value="Awaiting input...",
            elem_id="output-display"
        )

    with gr.Row():
        performance_graph = gr.Image(label="Performance Metrics Graph", visible=False)

    with gr.Row():
        loader = gr.Markdown(
            value="""
            <div class="loader"></div>
            """,
            visible=False,
            elem_id="loader-overlay"
        )

    with gr.Row():
        screenshot_button = gr.Button("Capture Screenshot", elem_id="screenshot-button")
        screenshot_output = gr.Image(label="Screenshot", visible=False)

    submit_button.click(
        fn=gradio_handler,
        inputs=[query_input],
        outputs=[output_display, metrics_display, performance_graph, loader],
        show_progress=False
    )

    def take_screenshot():
        screenshot = capture_screenshot(output_display.value)
        return gr.update(value=screenshot, visible=True)

    screenshot_button.click(
        fn=take_screenshot,
        inputs=[],
        outputs=[screenshot_output]
    )

# Launch Gradio interface
if __name__ == "__main__":
   interface.launch(share=True)

Context: None
Assistant: