from llama_cpp import Llama
import gradio as gr
def predict(message, history):
    # 1 & 2. Define chat function and construct prompt string including history
    prompt = ""
    for user_msg, assistant_msg in history:
        prompt += f"Q: {user_msg} A: {assistant_msg}\n"
    prompt += f"Q: {message} A:"

    # 3 & 4. Call llm with streaming and iterate over chunks
    # Using parameters from instructions: max_tokens and temperature
    response_iterator = llm(
        prompt,
        max_tokens=2048,
        stop=["Q:", "\n"],
        stream=True,
        temperature=0.7,
        top_p=0.95,
        top_k=50,
    )

    partial_message = ""
    for chunk in response_iterator:
        if "choices" in chunk and len(chunk["choices"]) > 0:
            text = chunk["choices"][0].get("text", "")
            partial_message += text
            yield partial_message

# 5. Create instance of gr.ChatInterface
demo = gr.ChatInterface(
    fn=predict,
    title="Llama 3.3 Thinking Heretic Chat",
    description="Interactive chat interface using llama-cpp-python and Gradio."
)

# 6. Execute launch method
demo.launch(share=True)