from llama_cpp import Llama
from duckduckgo_search import DDGS
import gradio as gr

# Load model
llm = Llama(model_path="models/Sam-reason-A1.Q4_K_S.gguf", n_ctx=2048)

# Define tools
def search_tool(query):
    with DDGS() as ddgs:
        results = ddgs.text(query)
        return "\n".join([r["title"] + ": " + r["href"] for r in results[:3]])

def calc_tool(expr):
    try: return str(eval(expr))
    except: return "Error in expression."

tools = {
    "search": search_tool,
    "calc": calc_tool
}

# Tool registry parser
def parse_tools(output):
    for name in tools:
        if f"<tool:{name}>" in output:
            arg_start = output.find(f"<tool:{name}>") + len(f"<tool:{name}>")
            arg_end = output.find(f"</tool:{name}>")
            arg = output[arg_start:arg_end].strip()
            result = tools[name](arg)
            return f"Tool [{name}] → {result}"
    return None

# Agent loop
def agent_chat(user_input, history=[]):
    history.append({"role": "user", "content": user_input})
    prompt = "\n".join([f"{m['role']}: {m['content']}" for m in history])

    output = llm(prompt=prompt, stop=["user:", "system:"], echo=False)
    response = output["choices"][0]["text"].strip()

    tool_result = parse_tools(response)
    if tool_result:
        response += f"\n{tool_result}"

    history.append({"role": "assistant", "content": response})
    return response

gr.ChatInterface(fn=agent_chat).launch()