import gradio as gr
from huggingface_hub import InferenceClient
from web import search  # Import web search tool

client = InferenceClient("Futuresony/future_ai_12_10_2024.gguf")

def is_uncertain(question, response):
    """Check if the model's response is unreliable."""
    # 1. If response length is too short, it's likely a guess.
    if len(response.split()) < 4:
        return True

    # 2. If response repeats the question, it might be unsure.
    if response.lower() in question.lower():
        return True
    
    # 3. If the response contains generic phrases like "Kulingana na utafiti" (According to research)
    uncertain_phrases = [
        "Kulingana na utafiti", "Inaaminika kuwa", "Ninadhani", 
        "It is believed that", "Some people say", "Inasemekana kuwa"
    ]
    if any(phrase.lower() in response.lower() for phrase in uncertain_phrases):
        return True
    
    return False

def google_search(query):
    """Fetch search results from Google."""
    results = search(query)
    if results:
        return results[0]  # Return the first result
    return "Sorry, I couldn't find an answer on Google."

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  # Stream partial responses
    
    # If response is unreliable, fetch from Google
    if is_uncertain(message, response):
        google_response = google_search(message)
        yield f"🤖 AI: {response}\n\n🌍 Google: {google_response}"

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()