app.py

import os
import gradio as gr
from bagoodex_client import BagoodexClient
from r_types import ChatMessage
from prompts import SYSTEM_PROMPT_FOLLOWUP, SYSTEM_PROMPT_MAP, SYSTEM_PROMPT_BASE
from helpers import (
    embed_video,
    embed_image,
    format_links,
    embed_google_map,
    format_knowledge,
    format_followup_questions
)

client = BagoodexClient()

# ----------------------------
# Chat & Follow-up Functions
# ----------------------------
def chat_function(message, history, followup_state, chat_history_state):
    """
    Process a new user message.
    Appends the message and response to the conversation,
    and retrieves follow-up questions.
    """
    # complete_chat returns a new followup id and answer
    followup_id_new, answer = client.complete_chat(message)
    # Update conversation history (if history is None, use an empty list)
    if history is None:
        history = []
    updated_history = history + [ChatMessage({"role": "user", "content": message}),
                                ChatMessage({"role": "assistant", "content": answer})]
    # Retrieve follow-up questions using the updated conversation
    followup_questions_raw = client.base_qna(
        messages=updated_history, system_prompt=SYSTEM_PROMPT_FOLLOWUP
    )
    # Format them using the helper
    followup_md = format_followup_questions(followup_questions_raw)
    return answer, followup_id_new, updated_history, followup_md

def handle_followup_click(question, followup_state, chat_history_state):
    """
    When a follow-up question is clicked, send it as a new message.
    """
    if not question:
        return chat_history_state, followup_state, ""
    # Process the follow-up question via complete_chat
    followup_id_new, answer = client.complete_chat(question)
    updated_history = chat_history_state + [ChatMessage({"role": "user", "content": question}),
                                            ChatMessage({"role": "assistant", "content": answer})]
    # Get new follow-up questions
    followup_questions_raw = client.base_qna(
        messages=updated_history, system_prompt=SYSTEM_PROMPT_FOLLOWUP
    )
    followup_md = format_followup_questions(followup_questions_raw)
    return updated_history, followup_id_new, followup_md

def handle_local_map_click(followup_state, chat_history_state):
    """
    On local map click, try to get a local map.
    If issues occur, fall back to using the SYSTEM_PROMPT_MAP.
    """
    if not followup_state:
        return chat_history_state
    try:
        result = client.get_local_map(followup_state)
        map_url = result.get('link', '')
        # Use helper to produce an embedded map iframe
        html = embed_google_map(map_url)
    except Exception:
        # Fall back: use the base_qna call with SYSTEM_PROMPT_MAP
        result = client.base_qna(
            messages=chat_history_state, system_prompt=SYSTEM_PROMPT_MAP
        )
        # Assume result contains a 'link' field
        html = embed_google_map(result.get('link', ''))
    new_message = ChatMessage({"role": "assistant", "content": html})
    return chat_history_state + [new_message]

def handle_knowledge_click(followup_state, chat_history_state):
    """
    On knowledge base click, fetch and format knowledge content.
    """
    if not followup_state:
        return chat_history_state
    result = client.get_knowledge(followup_state)
    md = format_knowledge(result)
    new_message = ChatMessage({"role": "assistant", "content": md})
    return chat_history_state + [new_message]

# ----------------------------
# Advanced Search Functions
# ----------------------------
def perform_image_search(followup_state):
    if not followup_state:
        return []
    result = client.get_images(followup_state)
    # For images we simply return a list of original URLs
    return [item.get("original", "") for item in result]

def perform_video_search(followup_state):
    if not followup_state:
        return "<p>No followup ID available.</p>"
    result = client.get_videos(followup_state)
    # Use the helper to produce the embed iframes (supports multiple videos)
    return embed_video(result)

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

# ----------------------------
# UI Build
# ----------------------------
css = """
#chatbot {
    height: 100%;
}
"""

with gr.Blocks(css=css, fill_height=True) as demo:
    # State variables to hold followup ID and conversation history, plus follow-up questions text
    followup_state = gr.State(None)
    chat_history_state = gr.State([])  # holds conversation history as a list of messages
    followup_md_state = gr.State("")     # holds follow-up questions as Markdown text

    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")
            # The ChatInterface now uses additional outputs for both followup_state and conversation history,
            # plus follow-up questions Markdown.
            chat = gr.ChatInterface(
                fn=chat_function,
                type="messages",
                additional_inputs=[followup_state, chat_history_state],
                additional_outputs=[followup_state, chat_history_state, followup_md_state],
            )
            # Button callbacks to append local map and knowledge base results to chat
            btn_local_map.click(
                fn=handle_local_map_click,
                inputs=[followup_state, chat_history_state],
                outputs=chat.chatbot
            )
            btn_knowledge.click(
                fn=handle_knowledge_click,
                inputs=[followup_state, chat_history_state],
                outputs=chat.chatbot
            )
            # Below the chat input, display follow-up questions and let user select one.
            followup_radio = gr.Radio(
                choices=[], label="Follow-up Questions (select one and click Send Follow-up)"
            )
            btn_send_followup = gr.Button("Send Follow-up")
            # When a follow-up question is sent, update the chat conversation, followup state, and follow-up list.
            btn_send_followup.click(
                fn=handle_followup_click,
                inputs=[followup_radio, followup_state, chat_history_state],
                outputs=[chat.chatbot, followup_state, followup_md_state]
            )
            # Also display the follow-up questions markdown (for reference) in a Markdown component.
            followup_markdown = gr.Markdown(label="Follow-up Questions", value="", visible=True)
            # When the followup_md_state updates, also update the radio choices.
            def update_followup_radio(md_text):
                # Assume the helper output is a Markdown string with list items.
                # We split the text to extract the question lines.
                lines = md_text.splitlines()
                questions = []
                for line in lines:
                    if line.startswith("- "):
                        questions.append(line[2:])
                return gr.update(choices=questions, value=None), md_text
            followup_md_state.change(
                fn=update_followup_radio,
                inputs=[followup_md_state],
                outputs=[followup_radio, followup_markdown]
            )
        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.HTML(label="Video Results")  # HTML for embedded video iframes
                links_output = gr.Markdown(label="Links Results")
                btn_images.click(
                    fn=perform_image_search,
                    inputs=[followup_state],
                    outputs=[gallery_output]
                )
                btn_videos.click(
                    fn=perform_video_search,
                    inputs=[followup_state],
                    outputs=[video_output]
                )
                btn_links.click(
                    fn=perform_links_search,
                    inputs=[followup_state],
                    outputs=[links_output]
                )
    demo.launch()