import gradio as gr
import requests
import os

# 環境変数からトークンを取得、なければデフォルト値を使用
token = os.environ.get("FRIENDLI_TOKEN") or "YOUR_FRIENDLI_TOKEN"

# APIエンドポイント
url = "https://api.friendli.ai/dedicated/v1/completions"

# ヘッダー情報
headers = {
    "Authorization": f"Bearer {token}",
    "Content-Type": "application/json"
}

def get_response(prompt, system_message, temperature, top_p, max_tokens):
    # ペイロードを作成
    payload = {
        "model": "jd5fsx11n7go",
        "prompt": f"{system_message}\n{prompt}",
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": top_p
    }
    
    # APIリクエストを送信
    response = requests.post(url, headers=headers, json=payload)
    
    # レスポンスをチェック
    if response.status_code == 200:
        return response.json().get("choices", [{}])[0].get("text", "回答が生成できませんでした。")
    else:
        return f"エラーが発生しました: {response.status_code} {response.text}"

def respond(message, history, system_message, max_tokens, temperature, top_p):
    # レスポンスを取得
    response = get_response(message, system_message, temperature, top_p, max_tokens)
    
    # ヒストリに追加
    history.append((message, "あなた"))
    history.append((response, "チャットボット"))
    
    return history, history

# Gradio ChatInterfaceの作成
demo = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        gr.Textbox(value="あなたはフレンドリーなチャットボットです。", label="システムメッセージ"),
        gr.Slider(minimum=1, maximum=2048, value=768, step=1, label="新規トークン最大"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="温度"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (核 sampling)"),
    ],
    title="Friendli AI チャットボット",
    description="Friendli AIを使用したチャットボットです。パラメータを調整して会話を行ってください。",
)

# インターフェースの起動
demo.launch(share=True)