import gradio as gr
from transformers import pipeline
from huggingface_hub import InferenceClient
import requests
from bs4 import BeautifulSoup
from Applio import *
# Initialize the text generation pipeline
#pipe = pipeline("text-generation", model="", trust_remote_code=True)
client = InferenceClient("mlabonne/Meta-Llama-3.1-8B-Instruct-abliterated")

def web_search(query):
    # Simulate a web search using Google
    response = requests.get(f"https://www.google.com/search?q={query}")
    soup = BeautifulSoup(response.text, "html.parser")
    results = []
    for g in soup.find_all('div', class_='BNeawe vvjwJb AP7Wnd'):
        results.append(g.get_text())
    return results

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]})

    # Check if message is a search request
    if "search:" in message.lower():
        search_query = message.split("search:", 1)[1].strip()
        search_results = web_search(search_query)
        response = "\n".join(search_results[:5])  # Return top 5 search results
    else:
        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

    yield response


demo = gr.ChatInterface(
    respond,
    title="INDONESIAN CHATBOT",
    theme=applio,
    additional_inputs=[
        gr.Textbox(value="Anda adalah Maya, sebuah asisten AI yang membantu. Anda menjawab pertanyaan pengguna seperti teman manusia. Anda juga Ahli di segala bidang dan juga belajar dan mencoba menjawab dari konteks terkait pertanyaan sebelumnya. Cobalah yang terbaik untuk memberikan respons terbaik kepada pengguna. Usahakan juga untuk menunjukkan emosi dan membalas seperti manusia, menggunakan bentuk yang singkat, nada dan emosi yang bersahabat dan respon menggunakan emoji sebagai emosi ", label="System message", visible=False),
        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.queue(api_open=False).launch(show_api=False)