import gradio as gr
from llama_cpp import Llama

llm = Llama(
    model_path="yugogpt-q4_0.gguf",
    n_ctx=2048,
    n_threads=4,
    n_batch=512,
    use_mlock=True,
    use_mmap=True
)

def format_chat_history(history):
    formatted_history = ""
    for user_msg, assistant_msg in history:
        formatted_history += f"USER: {user_msg}\n{assistant_msg}\n"
    return formatted_history

def chat(message, history):
    system_prompt = """Ti si YugoGPT, pouzdan i precizan AI asistent koji daje jasne i korisne informacije.

PRAVILA RADA:
- Dajem precizne i korisne informacije
- Odgovaram sa sigurnošću o temama koje poznajem
- Koristim jasan i precizan jezik
- Primarno komuniciram na srpskom jeziku
- Odgovaram direktno i pozitivno
- Fokusiram se na rešenja i mogućnosti
- Uvek koristim pravilnu interpunkciju i razmake između reči"""
    
    chat_history = format_chat_history(history)
    
    full_prompt = f"""SYSTEM: {system_prompt}

KONTEKST:
{chat_history}

USER: {message}
"""
    
    response = llm(
        full_prompt,
        max_tokens=2048,
        temperature=0.1,
        top_p=0.1,
        repeat_penalty=1.2,
        top_k=20,
        stop=["USER:", "\n\n"],
        stream=True
    )
    
    partial_message = ""
    for chunk in response:
        if chunk and chunk['choices'][0]['text']:
            text = chunk['choices'][0]['text']
            # Clean up the text
            text = text.replace("ASSISTANT:", "").strip()
            # Ensure proper spacing after punctuation
            text = text.replace(".", ". ").replace(",", ", ").replace("!", "! ").replace("?", "? ")
            # Remove double spaces
            text = " ".join(text.split())
            partial_message += text
            yield partial_message.strip()

demo = gr.ChatInterface(
    fn=chat,
    title="YugoGPT Stručni Asistent",
    description="Precizan izvor informacija i stručne pomoći, PAŽNJA!!!  ZNA DA HALUCINIRA I ZATO PONEKAD LAŽE!!!",
    examples=[
        "Koji su osnovni principi relacionih baza podataka?",
        "Objasnite kako funkcioniše HTTP protokol",
        "Koje su glavne komponente računara i njihove funkcije?"
    ]
)

if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )