import os
from openai import OpenAI
import gradio as gr

client = OpenAI(
    api_key=os.environ["GROQ_API_KEY"],
    base_url="https://api.groq.com/openai/v1",
)

class GroqLLM:
    def __init__(self, client, model_name):
        self.client = client
        self.model_name = model_name
        self.memory_history = []

    def predict(self, user_message):
        prompt = "You are a helpful assistant.\n"
        for msg in self.memory_history:
            prompt += f"{msg}\n"
        prompt += f"User: {user_message}\nAssistant:"

        response = self.client.responses.create(
            model=self.model_name,
            input=prompt,
            max_output_tokens=300
        )

        try:
            answer = response.output_text
        except:
            answer = str(response)

        self.memory_history.append(f"User: {user_message}")
        self.memory_history.append(f"Assistant: {answer}")
        if len(self.memory_history) > 20:
            self.memory_history = self.memory_history[-20:]
        return answer

llm = GroqLLM(client, "openai/gpt-oss-20b")

def chat_func(msg, history):
    return llm.predict(msg)

demo = gr.ChatInterface(chat_func)

demo.launch()