import gradio as gr
import requests
import json
from huggingface_hub import InferenceClient

API_TOKEN = "your_huggingface_api_token" # Replace with your actual token
API_URL = "https://api-inference.huggingface.co/models/InterSync/Mistral-7B-Instruct-v0.2-Function-Calling"
headers = {"Authorization": f"Bearer {API_TOKEN}"}

def get_weather(location: str, unit: str = "celsius"):
    # Replace with your actual weather API call
    pass 
    
def get_weather_schema():
    return {
        "name": "get_weather",
        "description": "Get the current weather in a given location",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {"type": "string", "description": "The city and state, or zip code"},
                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"], "description": "Unit of temperature"}
            },
            "required": ["location"]
        }
    }

def query_model(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()

with gr.Blocks() as demo:
    gr.Markdown("# Mistral-7B-Instruct Function Calling Demo")
    
    with gr.Row():
        with gr.Column(scale=4):
            input_text = gr.Textbox(label="Enter your text", lines=5)
            submit_btn = gr.Button("Submit")
        with gr.Column(scale=6):
            output_text = gr.Textbox(label="Model Output", lines=10)

    def user(user_message, history):
        return "", history + [[user_message, None]]  # Add user message to chat history

    def bot(history):
        if history:    
            user_message = history[-1][0]
            payload = {
                "inputs": user_message,
                "parameters": {"function_call": "auto"}
            }
            output = query_model(payload)
        else:
            return history  # Or some default response if history is empty
        # Parse the model's response
        if 'function_call' in output and 'name' in output['function_call']:
            function_name = output['function_call']['name']
            arguments = output['function_call'].get('arguments', {})
            if function_name == "get_weather" and arguments:
                weather_info = get_weather(**arguments)
                response_message = f"The weather in {arguments['location']} is {weather_info['description']} with a temperature of {weather_info['temperature']} {weather_info['unit']}."
            else:
                response_message = "Function not found or invalid arguments."
        else:
            response_message = output[0]['generated_text']

        history[-1][1] = response_message
        return history

    input_text.change(user, [input_text, output_text], [input_text, output_text], queue=False).then(
        bot, [output_text], [output_text]
    )
    submit_btn.click(user, [input_text, output_text], [input_text, output_text], queue=False).then(
        bot, [output_text], [output_text]
    )

demo.queue().launch()
"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

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

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response


For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)


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