import gradio as gr
import ollama  # Ensure ollama is available in your Space

# Define chat function
def chat_with_model(message, history):
    try:
        response = ollama.chat(
            model="mistral-small",
            messages=[{"role": "user", "content": message}]
        )
        return response['message']['content']
    except Exception as e:
        return f"Error: {str(e)}"

# Define Gradio Chat Interface
chatbot = gr.ChatInterface(
    chat_with_model,
    title="Mistral AI Chatbot",
    description="Chat with Mistral-Small using Ollama on Hugging Face Spaces.",
)

# Launch the Gradio app
if __name__ == "__main__":
    chatbot.launch()