from llama_cpp import Llama
import json
import time

MODEL_PATH = "./model.gguf"

llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=4096,
    n_gpu_layers=-1,
    verbose=False,
)

SYSTEM_PROMPT = """
You are a helpful AI assistant.
Answer clearly and accurately.
"""

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

print("=" * 60)
print("GGUF Agent Started")
print("Type 'exit' to quit")
print("=" * 60)

while True:

    user_input = input("\nUser: ")

    if user_input.lower() in ["exit", "quit"]:
        break

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

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

    assistant_text = response["choices"][0]["message"]["content"]

    print("\nAssistant:", assistant_text)

    messages.append(
        {
            "role": "assistant",
            "content": assistant_text
        }
    )