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

# Load GGUF model (Sam-reason-S2.1, gemma3 arch)
llm = Llama(
    model_path="models/mistral-7b-v0.1.Q4_K_S.gguf",
    n_ctx=2048,
    verbose=False
)

# Tools
def search_tool(q):
    with DDGS() as ddgs:
        results = ddgs.text(q)
        return "\n".join([f"{r['title']}: {r['href']}" for r in results[:3]])

def calc_tool(expr):
    try:
        return str(eval(expr))
    except Exception as e:
        return f"Math error: {e}"

def run_tool(command):
    with Sandbox(template="base") as sb:
        result = sb.run(command)
    return result.stdout or result.stderr or "No output."

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

# Tool parser
def parse_tools(text):
    for name in tools:
        if f"<tool:{name}>" in text:
            start = text.find(f"<tool:{name}>") + len(f"<tool:{name}>")
            end = text.find(f"</tool:{name}>")
            arg = text[start:end].strip()
            return tools[name](arg)
    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()
    result = parse_tools(response)
    if result:
        response += f"\n🔧 Tool Output:\n{result}"
    history.append({"role": "assistant", "content": response})
    return response

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