Update gradio_interface.py

gradio_interface.py

@@ -1,190 +1,34 @@
-# gradio_interface.py
-
 import gradio as gr
-from ai_model import query_ai_model, preprocess_chat_history
-import asyncio
-import logging
-from config import (
-    DEFAULT_MAX_OUTPUT_TOKENS,
-    DEFAULT_TEMPERATURE,
-    DEFAULT_TOP_P,
-    DEFAULT_TOP_K,
-)
-import os
-from typing import List, Optional  # <-- Added this import
-import time
-
-# Configure logging
-logging.basicConfig(level=logging.ERROR)
-
-def preprocess_stop_sequences(stop_sequences: str) -> Optional[List[str]]:
-    if not stop_sequences:
-        return None
-    return [sequence.strip() for sequence in stop_sequences.split(",")]
-
-def preprocess_chat_history(
-    history: List[tuple]
-) -> List[dict]:
-    messages = []
-    for user_message, model_message in history:
-        if user_message is not None:
-            messages.append({'role': 'user', 'content': user_message})
-        if model_message is not None:
-            messages.append({'role': 'assistant', 'content': model_message})
-    return messages
-
-def user(text_prompt: str, chatbot: List[tuple]):
-    if text_prompt:
-        chatbot.append((text_prompt, None))
-    return "", chatbot
-
-async def bot(
-    google_key: str,
-    temperature: float,
-    max_output_tokens: int,
-    stop_sequences: str,
-    top_k: int,
-    top_p: float,
-    chatbot: List[tuple]
-):
-    if len(chatbot) == 0:
-        yield chatbot
-        return
-
-    google_key = google_key if google_key else os.environ.get("GOOGLE_API_KEY")
-    if not google_key:
-        chatbot.append(('GOOGLE_API_KEY is not set. Please provide your API key.', None))
-        yield chatbot
-        return
-
-    genai.configure(api_key=google_key)
+from ai_model import query_ai_model  # Import AI model query function
 
-    stop_sequences_list = preprocess_stop_sequences(stop_sequences)
-
-    try:
-        # Process the latest user message
-        message = chatbot[-1][0]
-        history = preprocess_chat_history(chatbot[:-1])  # Exclude the last message
-
-        # Query the AI model
-        assistant_reply = await query_ai_model(
-            message=message,
-            history=history,
-            max_output_tokens=max_output_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            stop_sequences=stop_sequences_list
-        )
-
-        # Update the chat history with the assistant's reply
-        chatbot[-1] = (chatbot[-1][0], "")
-        response_content = assistant_reply['content']
-
-        # Streaming effect
-        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
-    except Exception as e:
-        logging.error("Error in bot function", exc_info=True)
-        error_message = f"An error occurred: {str(e)}"
-        chatbot[-1] = (chatbot[-1][0], error_message)
-        yield chatbot
+# Function to handle chatbot conversation
+def gradio_chatbot(text, history):
+    history.append((text, None))
+    response = asyncio.run(query_ai_model(text))
+    history[-1] = (text, response)
+    return history
 
+# Gradio interface
 def create_gradio_interface():
-    # 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
-    )
-    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."
-    )
-
+    chatbot = gr.Chatbot()
+    textbox = gr.Textbox(placeholder="Ask a question and press Enter", show_label=False)
+    
+    # Set up the Gradio interface layout
     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>")
-        with gr.Column():
-            google_key_component.render()
-            chatbot_component.render()
-            text_prompt_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,
-                temperature_component,
-                max_output_tokens_component,
-                stop_sequences_component,
-                top_k_component,
-                top_p_component,
-                chatbot_component
-            ],
-            outputs=chatbot_component,
-            queue=False
+        gr.HTML("<h1>Unified Gemini Chatbot</h1>")
+        chatbot.render()
+        textbox.render()
+
+        # Bind input textbox to chatbot response
+        textbox.submit(
+            fn=gradio_chatbot,
+            inputs=[textbox, chatbot],
+            outputs=chatbot
         )
 
     return demo
+
+# Launch Gradio interface
+def run_gradio_interface():
+    interface = create_gradio_interface()
+    interface.launch(share=True, debug=True)