from openai import OpenAI
import gradio as gr

YOUR_API_KEY = "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"

def predict(message, history):
    messages = [
        {
            "role": "system",
            "content": (
                "You are an artificial intelligence assistant and you need to "
                "engage in a helpful, detailed, polite conversation with a user."
            ),
        },
    ]

    for human, assistant in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": assistant})

    messages.append({"role": "user", "content": message})

    client = OpenAI(api_key=YOUR_API_KEY, base_url="https://api.perplexity.ai")

    response_stream = client.chat.completions.create(
        model="sonar-small-online",
        messages=messages,
        stream=True,
    )

    partial_message = ""
    for response in response_stream:
        if response.choices[0].finish_reason == "stop":
            partial_message += response.choices[0].text.strip()
            yield partial_message

gr.Interface(fn=predict, inputs="text", outputs="text").launch()