bhaiyasingh45/functiongemma-multiagent-router

pip install transformers torch

from transformers import AutoTokenizer, AutoModelForCausalLM
import re
import json

# Load model and tokenizer
model_name = "bhaiyahnsingh45/functiongemma-multiagent-router"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype="auto"
)

# Define your agent tools
from transformers.utils import get_json_schema

def technical_support_agent(issue_type: str, priority: str) -> str:
    """
    Routes technical issues to specialized support team.

    Args:
        issue_type: Type of technical issue (crash, authentication, performance, api_error, etc.)
        priority: Priority level (low, medium, high)
    """
    return f"Routing to Technical Support: {issue_type} with {priority} priority"

def billing_agent(request_type: str, urgency: str) -> str:
    """
    Routes billing and payment queries.

    Args:
        request_type: Type of request (refund, invoice, upgrade, cancellation, etc.)
        urgency: How urgent (low, medium, high)
    """
    return f"Routing to Billing: {request_type} with {urgency} urgency"

def product_info_agent(query_type: str, category: str) -> str:
    """
    Routes product information queries.

    Args:
        query_type: Type of query (features, comparison, integrations, limits, etc.)
        category: Category (plans, storage, mobile, security, etc.)
    """
    return f"Routing to Product Info: {query_type} about {category}"

# Get tool schemas
AGENT_TOOLS = [
    get_json_schema(technical_support_agent),
    get_json_schema(billing_agent),
    get_json_schema(product_info_agent)
]

# System message
SYSTEM_MSG = "You are an intelligent routing agent that directs customer queries to the appropriate specialized agent."

# Function to route queries
def route_query(user_query: str):
    """Route a user query to the appropriate agent"""

    messages = [
        {"role": "developer", "content": SYSTEM_MSG},
        {"role": "user", "content": user_query}
    ]

    # Format prompt
    inputs = tokenizer.apply_chat_template(
        messages,
        tools=AGENT_TOOLS,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt"
    )

    # Generate
    outputs = model.generate(
        **inputs.to(model.device),
        max_new_tokens=128,
        pad_token_id=tokenizer.eos_token_id
    )

    # Decode
    result = tokenizer.decode(
        outputs[0][len(inputs["input_ids"][0]):],
        skip_special_tokens=False
    )

    return result

# Example usage
query = "My app crashes when I try to upload large files"
result = route_query(query)
print(f"Query: {query}")
print(f"Routing: {result}")

<start_function_call>call:technical_support_agent{issue_type:crash,priority:high}<end_function_call>

query = "I'm getting a 500 error when calling the API"
result = route_query(query)
# Output: technical_support_agent(issue_type="api_error", priority="high")

query = "I need a refund for my annual subscription"
result = route_query(query)
# Output: billing_agent(request_type="refund", urgency="medium")

query = "What integrations do you support for project management?"
result = route_query(query)
# Output: product_info_agent(query_type="integrations", category="project_management")

def parse_function_call(output: str) -> dict:
    """Extract function name and arguments from model output"""

    pattern = r'<start_function_call>call:(\w+)\{([^}]+)\}<end_function_call>'
    match = re.search(pattern, output)

    if match:
        func_name = match.group(1)
        params_str = match.group(2)

        # Parse parameters
        params = {}
        param_pattern = r'(\w+):(?:<escape>(.*?)<escape>|([^,{}]+))'
        for p_match in re.finditer(param_pattern, params_str):
            key = p_match.group(1)
            val = p_match.group(2) or p_match.group(3).strip()
            params[key] = val

        return {
            "agent": func_name,
            "parameters": params
        }

    return {"agent": "unknown", "parameters": {}}

# Use it
query = "I was charged twice this month"
result = route_query(query)
parsed = parse_function_call(result)
print(parsed)
# Output: {'agent': 'billing_agent', 'parameters': {'request_type': 'dispute', 'urgency': 'high'}}

class MultiAgentRouter:
    def __init__(self, model_name: str):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            device_map="auto",
            torch_dtype="auto"
        )
        self.system_msg = "You are an intelligent routing agent..."

    def route(self, query: str) -> dict:
        """Route query and return agent + parameters"""
        messages = [
            {"role": "developer", "content": self.system_msg},
            {"role": "user", "content": query}
        ]

        inputs = self.tokenizer.apply_chat_template(
            messages,
            tools=AGENT_TOOLS,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt"
        )

        outputs = self.model.generate(
            **inputs.to(self.model.device),
            max_new_tokens=128,
            pad_token_id=self.tokenizer.eos_token_id
        )

        result = self.tokenizer.decode(
            outputs[0][len(inputs["input_ids"][0]):],
            skip_special_tokens=False
        )

        return parse_function_call(result)

# Usage
router = MultiAgentRouter("bhaiyahnsingh45/functiongemma-multiagent-router")
routing = router.route("My payment failed but I don't know why")
print(f"Route to: {routing['agent']}")
print(f"Parameters: {routing['parameters']}")

@misc{functiongemma_multiagent_router,
  author = {Bhaiya Singh},
  title = {Multi-Agent Router: Fine-tuned FunctionGemma for Customer Support},
  year = {2025},
  publisher = {Hugging Face},
  howpublished = {\url{https://huggingface.co/bhaiyahnsingh45/functiongemma-multiagent-router}}
}