import gradio as gr
from openai import OpenAI
import time
import re

# Available models
MODELS = [
    "Meta-Llama-3.1-405B-Instruct",
    "Meta-Llama-3.1-70B-Instruct",
    "Meta-Llama-3.1-8B-Instruct"
]

def create_client(api_key, base_url):
    return OpenAI(
        api_key=api_key,
        base_url=base_url
    )

def chat_with_ai(message, chat_history, system_prompt):
    messages = [
        {"role": "system", "content": system_prompt},
    ]
    
    for human, ai in chat_history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": ai})
    
    messages.append({"role": "user", "content": message})
    
    return messages

def respond(message, chat_history, model, system_prompt, thinking_budget, api_key, base_url):
    client = create_client(api_key, base_url)
    messages = chat_with_ai(message, chat_history, system_prompt.format(budget=thinking_budget))
    response = ""
    start_time = time.time()
    
    try:
        for chunk in client.chat.completions.create(
            model=model,
            messages=messages,
            stream=True
        ):
            content = chunk.choices[0].delta.content or ""
            response += content
            yield response, time.time() - start_time
    except Exception as e:
        yield f"Error: {str(e)}", time.time() - start_time

def parse_response(response):
    answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
    reflection_match = re.search(r'<reflection>(.*?)</reflection>', response, re.DOTALL)
    
    answer = answer_match.group(1).strip() if answer_match else ""
    reflection = reflection_match.group(1).strip() if reflection_match else ""
    
    steps = re.findall(r'<step>(.*?)</step>', response, re.DOTALL)
    
    return answer, reflection, steps

def process_chat(message, history, model, system_prompt, thinking_budget, api_key, base_url):
    if not api_key or not base_url:
        history.append((message, "Please provide both API Key and Base URL before starting the chat."))
        return history, history

    full_response = ""
    thinking_time = 0
    
    for response, elapsed_time in respond(message, history, model, system_prompt, thinking_budget, api_key, base_url):
        full_response = response
        thinking_time = elapsed_time
    
    if full_response.startswith("Error:"):
        history.append((message, full_response))
        return history, history

    answer, reflection, steps = parse_response(full_response)
    
    formatted_response = f"**Answer:** {answer}\n\n**Reflection:** {reflection}\n\n**Thinking Steps:**\n"
    for i, step in enumerate(steps, 1):
        formatted_response += f"**Step {i}:** {step}\n"
    
    formatted_response += f"\n**Thinking time:** {thinking_time:.2f} s"
    
    history.append((message, formatted_response))
    return history, history

with gr.Blocks() as demo:
    gr.Markdown("# Llama3.1-Instruct-O1")
    gr.Markdown("[Powered by Llama3.1 models through SN Cloud](https://sambanova.ai/fast-api?api_ref=907266)")
    
    with gr.Row():
        api_key = gr.Textbox(label="API Key", type="password")
        base_url = gr.Textbox(label="Base URL", value="https://api.endpoints.anyscale.com/v1")
    
    with gr.Row():
        model = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0])
        thinking_budget = gr.Slider(minimum=1, maximum=100, value=1, step=1, label="Thinking Budget")
    
    system_prompt = gr.Textbox(
        label="System Prompt",
        value="""
        You are a helpful assistant in normal conversation.
        When given a problem to solve, you are an expert problem-solving assistant. Your task is to provide a detailed, step-by-step solution to a given question. Follow these instructions carefully:

        1. Read the given question carefully and reset counter between <count> and </count> to {budget}
        2. Generate a detailed, logical step-by-step solution.
        3. Enclose each step of your solution within <step> and </step> tags.
        4. You are allowed to use at most {budget} steps (starting budget), keep track of it by counting down within tags <count> </count>, STOP GENERATING MORE STEPS when hitting 0, you don't have to use all of them.
        5. Do a self-reflection when you are unsure about how to proceed, based on the self-reflection and reward, decides whether you need to return to the previous steps.
        6. After completing the solution steps, reorganize and synthesize the steps into the final answer within <answer> and </answer> tags.
        7. Provide a critical, honest and subjective self-evaluation of your reasoning process within <reflection> and </reflection> tags.
        8. Assign a quality score to your solution as a float between 0.0 (lowest quality) and 1.0 (highest quality), enclosed in <reward> and </reward> tags.

        Example format:            
        <count> [starting budget] </count>
        
        <step> [Content of step 1] </step>
        <count> [remaining budget] </count>

        <step> [Content of step 2] </step>
        <reflection> [Evaluation of the steps so far] </reflection>
        <reward> [Float between 0.0 and 1.0] </reward>
        <count> [remaining budget] </count>

        <step> [Content of step 3 or Content of some previous step] </step>
        <count> [remaining budget] </count>
        
        ...
    
        <step>  [Content of final step] </step>
        <count> [remaining budget] </count>
        
        <answer> [Final Answer] </answer>

        <reflection> [Evaluation of the solution] </reflection>

        <reward> [Float between 0.0 and 1.0] </reward>
        """,
        lines=10
    )
    
    chatbot_ui = gr.Chatbot()
    msg = gr.Textbox(label="Type your message here...")
    clear = gr.Button("Clear Chat")

    chat_history = gr.State([])

    msg.submit(
        process_chat,  # Use the renamed function
        [msg, chat_history, model, system_prompt, thinking_budget, api_key, base_url],
        [chatbot_ui, chat_history]
    )
    clear.click(lambda: ([], []), None, [chatbot_ui, chat_history], queue=False)

demo.launch()