import gradio as gr
import requests
import os
import json

ACCESS_TOKEN = os.getenv("HF_TOKEN")

def respond(message, max_tokens=512, temperature=0.7, top_p=0.95):
    data = {
        "model": "command-r-plus:104b-fp16",
        "prompt": message,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": top_p
    }

    # API 요청
    response = requests.post("http://hugpu.ai:7877/api/generate", json=data)
    try:
        # 개별 JSON 객체들을 분리하여 처리
        responses = []
        for obj in response.text.strip().split('\n'):
            # 각 줄을 JSON으로 파싱
            result = json.loads(obj)
            if result.get("done", False):
                break
            responses.append(result.get('response', ''))
        
        # 결과 텍스트를 조합
        generated_text = ''.join(responses)
    except json.JSONDecodeError as e:
        print("Failed to decode JSON from response:", e)
        generated_text = "An error occurred while processing your request."
    
    return generated_text

demo = gr.Interface(
    fn=respond,
    inputs=[
        gr.Textbox(label="Your Message", placeholder="Type your message here..."),
        gr.Slider(minimum=1, maximum=2048, value=512, label="Max Tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-P")
    ],
    outputs=gr.Textbox(),
    title="Advanced AI Chatbot",
    description="Enter your message and receive a response generated by an advanced AI model.",
    theme="Nymbo/Nymbo_Theme"
)

if __name__ == "__main__":
    demo.launch()