using gr.ChatInterface

app.py

@@ -3,84 +3,85 @@ 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 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'
 
-# Bagoodex API client encapsulation
-class BagoodexClient:
-    def __init__(self, api_key=AIML_API_KEY, api_url=API_URL):
-        self.api_key = api_key
-        self.api_url = api_url
-        self.client = OpenAI(base_url=self.api_url, api_key=self.api_key)
-
-    def complete_chat(self, query):
-        response = self.client.chat.completions.create(
-            model="bagoodex/bagoodex-search-v1",
-            messages=[{"role": "user", "content": query}],
-        )
-        followup_id = response.id  # unique ID for followup searches
-        answer = response.choices[0].message.content
-        return followup_id, answer
-
-    def get_links(self, followup_id):
-        headers = {"Authorization": f"Bearer {self.api_key}"}
-        params = {"followup_id": followup_id}
-        response = requests.get(f"{self.api_url}/v1/bagoodex/links", headers=headers, params=params)
-        return response.json()
-
-    def get_images(self, followup_id):
-        headers = {"Authorization": f"Bearer {self.api_key}"}
-        params = {"followup_id": followup_id}
-        response = requests.get(f"{self.api_url}/v1/bagoodex/images", headers=headers, params=params)
-        return response.json()
-
-    def get_videos(self, followup_id):
-        headers = {"Authorization": f"Bearer {self.api_key}"}
-        params = {"followup_id": followup_id}
-        response = requests.get(f"{self.api_url}/v1/bagoodex/videos", headers=headers, params=params)
-        return response.json()
-
-    def get_local_map(self, followup_id):
-        headers = {"Authorization": f"Bearer {self.api_key}"}
-        params = {"followup_id": followup_id}
-        response = requests.get(f"{self.api_url}/v1/bagoodex/local-map", headers=headers, params=params)
-        return response.json()
+client = BagoodexClient()
 
-    def get_knowledge(self, followup_id):
-        headers = {"Authorization": f"Bearer {self.api_key}"}
-        params = {"followup_id": followup_id}
-        response = requests.get(f"{self.api_url}/v1/bagoodex/knowledge", headers=headers, params=params)
-        return response.json()
+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
 
-client = BagoodexClient()
+# Helper formatting functions
+def format_videos(result):
+    return [vid.get('link', '') for vid in result]
 
-# Callback functions
-def chat_response(user_input, chat_history, followup_id):
-    followup_id_new, answer = client.complete_chat(user_input)
-    chat_history = chat_history or []
-    chat_history.append((user_input, answer))
-    return "", chat_history, followup_id_new
+# 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 perform_local_map_search(followup_id, chat_history):
+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 chat_history
+        return chatbot_value
     result = client.get_local_map(followup_id)
     formatted = format_local_map(result)
-    chat_history.append((None, formatted))
-    return chat_history
+    new_message = {"role": "assistant", "content": formatted}
+    return chatbot_value + [new_message]
 
-def perform_knowledge_search(followup_id, chat_history):
+def append_knowledge(followup_id, chatbot_value):
     if not followup_id:
-        return chat_history
+        return chatbot_value
     result = client.get_knowledge(followup_id)
     formatted = format_knowledge(result)
-    chat_history.append((None, formatted))
-    return chat_history
+    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 []
@@ -88,19 +89,13 @@ def perform_image_search(followup_id):
     urls = format_images(result)
     return urls
 
-def perform_video_search(followup_id):
-    if not followup_id:
-        return "No followup ID available."
-    result = client.get_videos(followup_id)
-    return format_videos(result)
-
 def perform_links_search(followup_id):
     if not followup_id:
-        return "No followup ID available."
+        return gr.Markdown("No followup ID available.")
     result = client.get_links(followup_id)
     return format_links(result)
 
-# Custom CSS to help the chat component fill the available height
+# Custom CSS
 css = """
 #chatbot {
     height: 100%;
@@ -109,32 +104,29 @@ css = """
 
 # Build UI
 with gr.Blocks(css=css, fill_height=True) as demo:
+    followup_state = gr.State(None)
     with gr.Row():
-        # Left column: Chat interface
         with gr.Column(scale=3):
-            chatbot = gr.Chatbot(label="Bagoodex Chat", elem_id="chatbot")
-            # State variables: chat history and followup ID.
-            chat_history_state = gr.State([])
-            followup_state = gr.State(None)
-            # Two small buttons placed above the input field:
             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 input area: textbox and explicit submit button.
-            with gr.Row():
-                txt = gr.Textbox(placeholder="Enter your query here...", show_label=False, scale=7)
-                submit_btn = gr.Button("Submit", scale=1)
-            # Wire up the chat callbacks:
-            submit_btn.click(chat_response, inputs=[txt, chat_history_state, followup_state],
-                               outputs=[txt, chatbot, followup_state])
-            txt.submit(chat_response, inputs=[txt, chat_history_state, followup_state],
-                       outputs=[txt, chatbot, followup_state])
-            # The two additional search buttons update the chat history by appending a new assistant message.
-            btn_local_map.click(perform_local_map_search, inputs=[followup_state, chat_history_state],
-                                outputs=[chatbot, chat_history_state])
-            btn_knowledge.click(perform_knowledge_search, inputs=[followup_state, chat_history_state],
-                                outputs=[chatbot, chat_history_state])
-        # Right column: Advanced search options for images, videos, and links.
+            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"):
@@ -142,9 +134,21 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                 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")
-                links_output = gr.HTML(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()
+                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()