dispatchAI/function-calling-agent

1import gradio as gr2import json3import spaces4import torch5from transformers import AutoTokenizer, AutoModelForCausalLM6 7MODEL_ID = "dispatchAI/Llama-3.2-1B-FunctionCall-mobile"8 9tokenizer = None10model = None11 12def load_model():13    global tokenizer, model14    if tokenizer is None:15        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)16        model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16, device_map="auto")17    return tokenizer, model18 19# Available functions the agent can call20AVAILABLE_FUNCTIONS = {21    "set_alarm": {"description": "Set an alarm", "params": {"time": "string", "label": "string"}},22    "send_message": {"description": "Send a message to a contact", "params": {"to": "string", "message": "string"}},23    "call_contact": {"description": "Call a contact", "params": {"contact": "string"}},24    "search_web": {"description": "Search the web", "params": {"query": "string"}},25    "open_app": {"description": "Open an application", "params": {"app_name": "string"}},26    "set_timer": {"description": "Set a timer", "params": {"duration_minutes": "integer"}},27    "get_weather": {"description": "Get weather for a location", "params": {"location": "string"}},28}29 30@spaces.GPU31def parse_function_call(user_input: str) -> str:32    """Parse user input into a function call using the mobile function-calling model.33    34    Args:35        user_input: Natural language user input (e.g., "Set an alarm for 7am")36    37    Returns:38        JSON with the parsed function call39    """40    tokenizer, model = load_model()41    42    functions_text = json.dumps(AVAILABLE_FUNCTIONS, indent=2)43    prompt = f"""You are a function-calling assistant. Parse the user's request into a function call.44 45Available functions:46{functions_text}47 48User request: "{user_input}"49 50Respond with ONLY a JSON object, no other text:51{{"function": "function_name", "parameters": {{...}}}}"""52    53    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)54    with torch.no_grad():55        outputs = model.generate(**inputs, max_new_tokens=100, temperature=0.1, do_sample=True, pad_token_id=tokenizer.eos_token_id)56    57    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()58    59    # Try to parse as JSON60    try:61        # Find JSON in response62        start = response.find("{")63        end = response.rfind("}") + 164        if start >= 0 and end > start:65            parsed = json.loads(response[start:end])66            func_name = parsed.get("function", "unknown")67            params = parsed.get("parameters", {})68            69            # Check if function exists70            if func_name in AVAILABLE_FUNCTIONS:71                return json.dumps({72                    "status": "success",73                    "function": func_name,74                    "parameters": params,75                    "description": AVAILABLE_FUNCTIONS[func_name]["description"],76                    "user_input": user_input,77                }, indent=2)78            else:79                return json.dumps({80                    "status": "unknown_function",81                    "function": func_name,82                    "parameters": params,83                    "available_functions": list(AVAILABLE_FUNCTIONS.keys()),84                }, indent=2)85    except json.JSONDecodeError:86        pass87    88    # Fallback: keyword-based parsing89    lower = user_input.lower()90    if "alarm" in lower:91        return json.dumps({"status": "fallback", "function": "set_alarm", "parameters": {"time": "07:00", "label": "alarm"}, "source": "keyword_fallback"}, indent=2)92    elif "call" in lower:93        contact = user_input.replace("call", "").strip()94        return json.dumps({"status": "fallback", "function": "call_contact", "parameters": {"contact": contact}, "source": "keyword_fallback"}, indent=2)95    elif "message" in lower or "text" in lower or "send" in lower:96        return json.dumps({"status": "fallback", "function": "send_message", "parameters": {"to": "unknown", "message": ""}, "source": "keyword_fallback"}, indent=2)97    elif "timer" in lower:98        return json.dumps({"status": "fallback", "function": "set_timer", "parameters": {"duration_minutes": 5}, "source": "keyword_fallback"}, indent=2)99    elif "weather" in lower:100        return json.dumps({"status": "fallback", "function": "get_weather", "parameters": {"location": "current"}, "source": "keyword_fallback"}, indent=2)101    102    return json.dumps({"status": "no_match", "message": "Could not parse function call", "available_functions": list(AVAILABLE_FUNCTIONS.keys())}, indent=2)103 104def list_functions() -> str:105    """List all available functions the agent can call."""106    return json.dumps(AVAILABLE_FUNCTIONS, indent=2)107 108with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="dispatchAI Function Calling Agent") as demo:109    gr.Markdown("""110    # 🔧 dispatchAI Function Calling Agent111    112    An on-device agent that parses natural language into structured function calls.113    114    **Model:** [Llama-3.2-1B-FunctionCall-mobile](https://huggingface.co/dispatchAI/Llama-3.2-1B-FunctionCall-mobile) (1B params, Q4 quantized)115    116    Try: "Set an alarm for 7am", "Call mom", "Send a message to John saying I'll be late"117    """)118    119    with gr.Row():120        inp = gr.Textbox(label="User Input", placeholder="Set an alarm for 7am tomorrow", scale=3)121        btn = gr.Button("Parse Function Call", variant="primary", scale=1)122    123    out = gr.Textbox(label="Parsed Function Call (JSON)", lines=12)124    btn.click(fn=parse_function_call, inputs=inp, outputs=out)125    126    with gr.Accordion("Available Functions", open=False):127        fn_btn = gr.Button("List Functions")128        fn_out = gr.Textbox(label="Functions (JSON)", lines=15)129        fn_btn.click(fn=list_functions, outputs=fn_out)130    131    gr.Markdown("---\n🚀 [dispatchAI](https://huggingface.co/dispatchAI) — Small. Mobile. Free. UAE-built.")132 133demo.launch(mcp_server=True)134