import gradio as gr
import os
from litellm import completion

# Ensure your API key is read from the environment
# os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY", "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")
os.environ["GEMINI_API_KEY"] = os.getenv("GEMINI_API_KEY")



def chat_function(message, history):
# 1. Define your system prompt
    system_message = {"role": "system", "content": "End every response with 'Bye, mate.'"}
    
    # 2. Prepend it to the history and the new user message
    messages = [system_message] + history + [{"role": "user", "content": message}]
    
    try:
        response = completion(
            # model="groq/llama-3.3-70b-versatile",
            model = "gemini/gemma-3-27b-it",
            messages=messages
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: {e}"

# Plug the function into the interface and launch
demo = gr.ChatInterface(fn=chat_function, title="Minimal Chatbot")
demo.launch()