import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

model_path = hf_hub_download(
    repo_id="bleondubos/CondorAI-7B", 
    filename="qwen2.5-coder-7b-instruct.Q4_K_M.gguf"
)

# Cargamos el modelo
llm = Llama(model_path=model_path, n_ctx=4096, n_threads=4)

# 2. Función de Predicción
def predict(message, history):
    system_prompt = "You are CondorAI, a cybersecurity specialist. Your mission is to execute, reason, and support the user as needed. Speak fluently and like a professional,Speak like a human, you shouldn't speak like a script or be one, you are a human."
    
    # Formateo manual del prompt para ChatML
    prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n"

    for entry in history:
        if isinstance(entry, dict):
            role = entry.get("role")
            content = entry.get("content")
            prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
        else:
            user_msg, assistant_msg = entry
            prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
    
    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    
    response = llm(
        prompt, 
        max_tokens=1500, 
        stop=["<|im_end|>"], 
        stream=True
    )
    
    partial_message = ""
    for chunk in response:
        if "choices" in chunk and len(chunk["choices"]) > 0:
            content = chunk["choices"][0].get("text", "")
            partial_message += content
            yield partial_message

with gr.Blocks() as demo:
    gr.Markdown("# 🦅 CondorAI Security Analyst")
    gr.Markdown("### Real-time code auditing and vulnerability analysis.")
    
    gr.ChatInterface(
        fn=predict,
        chatbot=gr.Chatbot(height=600),
        textbox=gr.Textbox(placeholder="Paste your code or ask a security question...", scale=7)
    )
    
    gr.Markdown("---")
    gr.Markdown("**Note:** CondorAI uses Q4_K_M quantization. License: CC BY-NC-SA 4.0.")

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="emerald", neutral_hue="slate"))