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 "

No followup ID available.

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