app.py

import os
import requests
import gradio as gr
from openai import OpenAI
from dotenv import load_dotenv
# from helpers import format_images, format_videos, format_links, format_local_map, format_knowledge
from bagoodex_client import BagoodexClient

# Load API key from .env
load_dotenv()
AIML_API_KEY = os.getenv('AIML_API_KEY')
API_URL = 'https://api.aimlapi.com'

client = BagoodexClient()

def format_knowledge(result):
    title = result.get('title', 'Unknown')
    type_ = result.get('type', '')
    born = result.get('born', '')
    died = result.get('died', '')
    content = f"""
    **{title}**  
    Type: {type_}  
    Born: {born}  
    Died: {died}
    """
    return gr.Markdown(content)

def format_images(result):
    urls = [item.get("original", "") for item in result]
    return urls

# Helper formatting functions
def format_videos(result):
    return [vid.get('link', '') for vid in result]

# Advanced search functions
def perform_video_search(followup_id):
    if not followup_id:
        return []
    result = client.get_videos(followup_id)
    return format_videos(result)

def format_links(result):
    links_md = "**Links:**\n"
    for url in result:
        title = url.rstrip('/').split('/')[-1]
        links_md += f"- [{title}]({url})\n"
    return gr.Markdown(links_md)

# Define the chat function
def chat_function(message, history, followup_id):
    followup_id_new, answer = client.complete_chat(message)
    return answer, followup_id_new

def format_local_map(result):
    link = result.get('link', '')
    image_url = result.get('image', '')
    html = f"""
    <div>
        <strong>Local Map:</strong><br>
        <a href='{link}' target='_blank'>View on Google Maps</a><br>
        <img src='{image_url}' style='width:100%;'/>
    </div>
    """
    return gr.HTML(html)

def append_local_map(followup_id, chatbot_value):
    if not followup_id:
        return chatbot_value
    result = client.get_local_map(followup_id)
    formatted = format_local_map(result)
    new_message = {"role": "assistant", "content": formatted}
    return chatbot_value + [new_message]

def append_knowledge(followup_id, chatbot_value):
    if not followup_id:
        return chatbot_value
    result = client.get_knowledge(followup_id)
    formatted = format_knowledge(result)
    new_message = {"role": "assistant", "content": formatted}
    return chatbot_value + [new_message]

# Define advanced search functions
def perform_image_search(followup_id):
    if not followup_id:
        return []
    result = client.get_images(followup_id)
    urls = format_images(result)
    return urls

def perform_links_search(followup_id):
    if not followup_id:
        return gr.Markdown("No followup ID available.")
    result = client.get_links(followup_id)
    return format_links(result)

# Custom CSS
css = """
#chatbot {
    height: 100%;
}
"""

# Build UI
with gr.Blocks(css=css, fill_height=True) as demo:
    followup_state = gr.State(None)
    with gr.Row():
        with gr.Column(scale=3):
            with gr.Row():
                btn_local_map = gr.Button("Local Map Search", variant="secondary", size="sm")
                btn_knowledge = gr.Button("Knowledge Base", variant="secondary", size="sm")
            chat = gr.ChatInterface(
                fn=chat_function,
                type="messages",
                additional_inputs=[followup_state],
                additional_outputs=[followup_state],
            )
            # Wire up the buttons to append to chat history
            btn_local_map.click(
                append_local_map,
                inputs=[followup_state, chat.chatbot],
                outputs=chat.chatbot
            )
            btn_knowledge.click(
                append_knowledge,
                inputs=[followup_state, chat.chatbot],
                outputs=chat.chatbot
            )
        with gr.Column(scale=1):
            gr.Markdown("### Advanced Search Options")
            with gr.Column(variant="panel"):
                btn_images = gr.Button("Search Images")
                btn_videos = gr.Button("Search Videos")
                btn_links = gr.Button("Search Links")
                gallery_output = gr.Gallery(label="Image Results", columns=2)
                video_output = gr.Gallery(label="Video Results", columns=1, visible=True)
                links_output = gr.Markdown(label="Links Results")
                btn_images.click(
                    perform_image_search,
                    inputs=[followup_state],
                    outputs=[gallery_output]
                )
                btn_videos.click(
                    perform_video_search,
                    inputs=[followup_state],
                    outputs=[video_output]
                )
                btn_links.click(
                    perform_links_search,
                    inputs=[followup_state],
                    outputs=[links_output]
                )
    demo.launch()