Update gradio_interface.py

gradio_interface.py

@@ -1,5 +1,7 @@
+# gradio_interface.py
+
 import gradio as gr
-from ai_model import query_ai_model
+from ai_model import query_ai_model, query_ai_model_vision, preprocess_chat_history
 import asyncio
 import logging
 from config import (
@@ -9,7 +11,9 @@ from config import (
     DEFAULT_TOP_K,
 )
 import os
-from typing import List, Optional
+from typing import List, Dict, Union, Optional
+from PIL import Image
+import time
 
 # Configure logging
 logging.basicConfig(level=logging.ERROR)
@@ -19,89 +23,201 @@ def preprocess_stop_sequences(stop_sequences: str) -> Optional[List[str]]:
         return None
     return [sequence.strip() for sequence in stop_sequences.split(",")]
 
-def submit_message(message, chat_history):
-    chat_history = chat_history or []
-    chat_history.append({'role': 'user', 'content': message})
-    return "", chat_history
-
-async def gradio_chatbot(
-    message, chat_history, google_key, temperature, max_output_tokens,
-    stop_sequences, top_k, top_p
+def user(text_prompt: str, chatbot: List[tuple]):
+    if text_prompt:
+        chatbot.append((text_prompt, None))
+    return "", chatbot
+
+async def bot(
+    google_key: str,
+    files: Optional[List[gr.File]],
+    temperature: float,
+    max_output_tokens: int,
+    stop_sequences: str,
+    top_k: int,
+    top_p: float,
+    chatbot: List[tuple]
 ):
+    if len(chatbot) == 0:
+        return chatbot
+
     google_key = google_key if google_key else os.environ.get("GOOGLE_API_KEY")
     if not google_key:
-        chat_history.append({'role': 'assistant', 'content': 'GOOGLE_API_KEY is not set. Please provide your API key.'})
-        return chat_history
+        raise ValueError(
+            "GOOGLE_API_KEY is not set. "
+            "Please follow the instructions in the README to set it up."
+        )
 
     genai.configure(api_key=google_key)
-
-    stop_sequences_list = preprocess_stop_sequences(stop_sequences)
-
-    try:
-        # Query the AI model
-        assistant_reply = await query_ai_model(
-            message=chat_history[-1]['content'],  # Last user message
-            history=chat_history[:-1],            # Exclude the last message
+    generation_config = genai.types.GenerationConfig(
+        temperature=temperature,
+        max_output_tokens=max_output_tokens,
+        stop_sequences=preprocess_stop_sequences(stop_sequences=stop_sequences),
+        top_k=top_k,
+        top_p=top_p
+    )
+
+    if files:
+        text_prompt = [chatbot[-1][0]] if chatbot[-1][0] and isinstance(chatbot[-1][0], str) else []
+        image_prompt = [Image.open(file.name).convert('RGB') for file in files]
+        response = await query_ai_model_vision(
+            text_prompt=text_prompt,
+            image_prompt=image_prompt,
             max_output_tokens=max_output_tokens,
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
-            stop_sequences=stop_sequences_list,
+            stop_sequences=preprocess_stop_sequences(stop_sequences=stop_sequences)
+        )
+    else:
+        messages = preprocess_chat_history(chatbot)
+        response = await query_ai_model(
+            message=chatbot[-1][0],
+            history=messages[:-1],  # Exclude the last user message
+            max_output_tokens=max_output_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            stop_sequences=preprocess_stop_sequences(stop_sequences=stop_sequences)
         )
 
-        # Append assistant's reply to the history
-        chat_history.append(assistant_reply)
-    except Exception as e:
-        logging.error("Error in gradio_chatbot", exc_info=True)
-        error_message = f"An error occurred: {str(e)}"
-        chat_history.append({'role': 'assistant', 'content': error_message})
-
-    return chat_history
+    # Streaming effect
+    chatbot[-1] = (chatbot[-1][0], "")
+    if response['content'].startswith("An error occurred"):
+        chatbot[-1] = (chatbot[-1][0], response['content'])
+        yield chatbot
+    else:
+        for i in range(0, len(response['content']), 10):
+            section = response['content'][i:i + 10]
+            chatbot[-1] = (chatbot[-1][0], chatbot[-1][1] + section)
+            await asyncio.sleep(0.01)
+            yield chatbot
 
 def create_gradio_interface():
-    with gr.Blocks() as demo:
-        gr.HTML("<h1 align='center'>Gemini Chat Interface 💬</h1>")
-        gr.HTML("<h2 align='center'>Interact with the Gemini AI model.</h2>")
+    # Define all components first
+    google_key_component = gr.Textbox(
+        label="GOOGLE API KEY",
+        value="",
+        type="password",
+        placeholder="Enter your Google API key",
+        info="You have to provide your own GOOGLE_API_KEY for this app to function properly",
+        visible=os.environ.get("GOOGLE_API_KEY") is None
+    )
+    chatbot_component = gr.Chatbot(
+        label='Gemini Chatbot',
+        height=600,
+        type='messages'
+    )
+    text_prompt_component = gr.Textbox(
+        placeholder="Hi there! [press Enter]", show_label=False
+    )
+    upload_button_component = gr.UploadButton(
+        label="Upload Images", file_count="multiple", file_types=["image"]
+    )
+    temperature_component = gr.Slider(
+        minimum=0,
+        maximum=1.0,
+        value=DEFAULT_TEMPERATURE,
+        step=0.05,
+        label="Temperature",
+        info="Controls the randomness of the AI's response."
+    )
+    max_output_tokens_component = gr.Slider(
+        minimum=1,
+        maximum=2048,
+        value=DEFAULT_MAX_OUTPUT_TOKENS,
+        step=1,
+        label="Max Output Tokens",
+        info="Controls the length of the AI's response."
+    )
+    stop_sequences_component = gr.Textbox(
+        label="Add stop sequence",
+        value="",
+        type="text",
+        placeholder="STOP, END",
+        info="Stops the AI response when any of these sequences are generated."
+    )
+    top_k_component = gr.Slider(
+        minimum=1,
+        maximum=40,
+        value=DEFAULT_TOP_K,
+        step=1,
+        label="Top-K",
+        info="Limits the next token selection to the K most probable tokens."
+    )
+    top_p_component = gr.Slider(
+        minimum=0,
+        maximum=1,
+        value=DEFAULT_TOP_P,
+        step=0.01,
+        label="Top-P",
+        info="Limits the next token selection to tokens with cumulative probability up to P."
+    )
 
-        google_key_component = gr.Textbox(
-            label="GOOGLE API KEY",
-            value="",
-            type="password",
-            placeholder="Enter your Google API key",
-            visible=os.environ.get("GOOGLE_API_KEY") is None
-        )
-
-        chatbot_component = gr.Chatbot(
-            label='Gemini Chatbot',
-            height=600,
-            type='messages'
+    with gr.Blocks() as demo:
+        gr.HTML("<h1 align='center'>Gemini Playground 💬</h1>")
+        gr.HTML("<h2 align='center'>Play with Gemini Pro and Gemini Pro Vision API</h2>")
+        gr.HTML("""
+        <div style="text-align: center; display: flex; justify-content: center; align-items: center;">
+            <a href="https://huggingface.co/spaces/SkalskiP/ChatGemini?duplicate=true">
+                <img src="https://bit.ly/3gLdBN6" alt="Duplicate Space" style="margin-right: 10px;">
+            </a>
+            <span>Duplicate the Space and run securely with your 
+                <a href="https://makersuite.google.com/app/apikey">GOOGLE API KEY</a>.
+            </span>
+        </div>
+        """)
+        with gr.Column():
+            google_key_component.render()
+            chatbot_component.render()
+            with gr.Row():
+                text_prompt_component.render()
+                upload_button_component.render()
+            with gr.Accordion("Parameters", open=False):
+                temperature_component.render()
+                max_output_tokens_component.render()
+                stop_sequences_component.render()
+                with gr.Accordion("Advanced", open=False):
+                    top_k_component.render()
+                    top_p_component.render()
+
+        # Connect the user input submission
+        text_prompt_component.submit(
+            fn=user,
+            inputs=[text_prompt_component, chatbot_component],
+            outputs=[text_prompt_component, chatbot_component],
+            queue=False
+        ).then(
+            fn=bot,
+            inputs=[
+                google_key_component,
+                upload_button_component,
+                temperature_component,
+                max_output_tokens_component,
+                stop_sequences_component,
+                top_k_component,
+                top_p_component,
+                chatbot_component
+            ],
+            outputs=chatbot_component,
+            queue=False
         )
 
-        with gr.Row():
-            msg = gr.Textbox(
-                label="Your Message",
-                placeholder="Type your message here and press Enter"
-            )
-
-        # Advanced parameters components...
-
-        msg.submit(
-            fn=submit_message,
-            inputs=[msg, chatbot_component],
-            outputs=[msg, chatbot_component]
-        ).then(
-            fn=gradio_chatbot,
+        # Handle image uploads
+        upload_button_component.upload(
+            fn=bot,
             inputs=[
-                msg,
-                chatbot_component,
                 google_key_component,
+                upload_button_component,
                 temperature_component,
                 max_output_tokens_component,
                 stop_sequences_component,
                 top_k_component,
-                top_p_component
+                top_p_component,
+                chatbot_component
             ],
-            outputs=chatbot_component
+            outputs=chatbot_component,
+            queue=False
         )
 
     return demo