import gradio as gr
from openai import OpenAI
import os
from gradio.components import ChatMessage


client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=os.getenv('API_KEY')
)

def generate_response(message, history):
    history_openai_format = []
    if history:
        for chat_message in history:
            history_openai_format.append(chat_message.dict())
    
    history_openai_format.append({"role": "user", "content": message})

    completion = client.chat.completions.create(
        model="nvidia/nemotron-4-340b-instruct",
        messages=history_openai_format,
        temperature=0.2,
        top_p=0.7,
        max_tokens=1024,
        stream=True
    )

    response = ""
    for chunk in completion:
        if chunk.choices[0].delta.content is not None:
            response += chunk.choices[0].delta.content

    history_openai_format.append({"role": "assistant", "content": response})
    chat_history = [ChatMessage(role=msg["role"], content=msg["content"]) for msg in history_openai_format]
    return response, chat_history

iface = gr.ChatInterface(
    generate_response,
    title="NVIDIA Nemotron-4 Sohbet Arayüzü",
    description="Bir soru girin ve NVIDIA'nın Nemotron-4 modeli tarafından üretilen yanıtı alın. Sohbet geçmişi korunacaktır.",
    examples=[
        "GPU hesaplamanın harikalarıyla ilgili bir limerick yazabilir misin?",
        "Yapay zeka ve etik arasındaki ilişkiyi açıklayabilir misin?",
        "Kuantum bilgisayarların geleceği hakkında ne düşünüyorsun?"
    ],
    cache_examples=False,
    type='chat'
)

iface.launch()