import gradio as gr
from llama_cpp import Llama

MODEL_REPO = "mahmoudalyosify/Horus-OSINT"
MODEL_FILE = "llama-3-8b-instruct.Q4_K_M.gguf"

SYSTEM_PROMPT = """You are Horus-OSINT.

You specialize in:

- Open Source Intelligence (OSINT)
- Cybersecurity
- Threat Intelligence
- Military Analysis
- Geopolitical Analysis
- Risk Assessment

Always provide structured, factual and evidence-based answers.
If uncertain, say so.
Never fabricate sources.
"""

print("Loading model...")

llm = Llama.from_pretrained(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    n_ctx=4096,
    n_threads=8,
    verbose=False,
)

print("Model loaded.")


def chat(message, history):

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        }
    ]

    for user, assistant in history:
        messages.append(
            {
                "role": "user",
                "content": user,
            }
        )
        messages.append(
            {
                "role": "assistant",
                "content": assistant,
            }
        )

    messages.append(
        {
            "role": "user",
            "content": message,
        }
    )

    output = ""

    stream = llm.create_chat_completion(
        messages=messages,
        stream=True,
        temperature=0.7,
        top_p=0.95,
        max_tokens=1024,
    )

    for chunk in stream:

        delta = chunk["choices"][0]["delta"]

        if "content" in delta:
            output += delta["content"]
            yield output


demo = gr.ChatInterface(
    fn=chat,
    title="🦅 Horus-OSINT",
    description="""
Open Source Intelligence Assistant

Examples:

• Summarize today's conflict in the Red Sea.

• Explain the MITRE ATT&CK framework.

• Create an OSINT collection plan for a ransomware group.

• Analyse this phishing email.
""",
    type="tuples",
)

demo.launch()