Here's a jinja2 template for function call (tested with llama-server)

#16
by tarruda - opened

I adapted from the vLLM template: https://github.com/vllm-project/vllm/pull/17149/files

Refactored it and changed it for the JSON style function call output as explained in https://ai.google.dev/gemma/docs/capabilities/function-calling

{#- Begin-of-sequence token -#}
{{ bos_token }}

{#- Extract system message if present -#}
{%- set system_message = "" -%}
{%- set start_index = 0 -%}
{%- if messages[0]['role'] == 'system' -%}
    {%- if messages[0]['content'] is string -%}
        {%- set system_message = messages[0]['content'] + '\n\n' -%}
    {%- else -%}
        {%- set system_message = messages[0]['content'][0]['text'] + '\n\n' -%}
    {%- endif -%}
    {%- set start_index = 1 -%}
{%- endif -%}

{#- Set tools to none if not defined -#}
{%- set tools = tools | default(none) -%}

{#- Validate message alternation (user/assistant pattern) -#}
{%- set ns = namespace(last_role='assistant') -%}
{%- for message in messages[start_index:] -%}

    {%- set role = message['role'] -%}
    {%- if message['role'] in ['user', 'tool'] -%}
        {%- set role = 'user' -%}
    {%- elif message['role'] == 'assistant' or 'tool_calls' in message -%}
        {%- set role = 'assistant' -%}
    {%- endif -%}

    {%- if role == ns.last_role -%}
        {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
    {%- endif -%}
    {%- set ns.last_role = role -%}
{%- endfor -%}

{#- Process each message -#}
{%- for message in messages[start_index:] -%}
    {#- Map roles to Gemma format -#}
    {%- if message['role'] == 'assistant' -%}
        {%- set role = 'model' -%}
    {%- elif message['role'] == 'tool' -%}
        {%- set role = 'user' -%}
    {%- else -%}
        {%- set role = message['role'] -%}
    {%- endif -%}
    
    {#- Start message turn -#}
    {{ '<start_of_turn>' + role + '\n' -}}
    
    {#- Include system message and tool definitions for first user message -#}
    {%- if loop.first and role == 'user' -%}
        {{ system_message }}
        {%- if tools is not none -%}
            {{ "You have access to functions. If you decide to invoke any of the function(s), you MUST put it in the format of\n" -}}
            {{ '{"name": function name, "parameters": dictionary of argument name and its value}\n\n' -}}
            {{ "You SHOULD NOT include any other text in the response if you call a function.\n" -}}
            {{ "If none of the functions can be used, point it out. If you lack the parameters required by the function, also point it out.\n" -}}
            {{ "Here is a list of functions available:\n" -}}
            {{ tools | tojson(indent=2) }}
            {{ "\n\n" -}}
        {%- endif -%}
    {%- endif -%}
    
    {#- Handle tool calls from assistant - format as JSON -#}
    {%- if 'tool_calls' in message -%}
        {%- for tool_call in message.tool_calls -%}
            {#- Extract function details -#}
            {%- set func = tool_call.function if tool_call.function is defined else tool_call -%}
            {#- Create JSON object for function call -#}
            {%- set call_json = {"name": func.name, "parameters": func.arguments} -%}
            {{ call_json | tojson }}
            {%- if not loop.last -%}{{ '\n' }}{%- endif -%}
        {%- endfor -%}
    {%- endif -%}
    
    {#- Handle tool response wrapper -#}
    {%- if message['role'] == 'tool' -%}
        {{ '<tool_response>\n' -}}
    {%- endif -%}
    
    {#- Render message content -#}
    {%- if message['content'] is string -%}
        {{ message['content'] | trim }}
    {%- elif message['content'] is iterable -%}
        {%- for item in message['content'] -%}
            {%- if item['type'] == 'image' -%}
                {{ '<start_of_image>' }}
            {%- elif item['type'] == 'text' -%}
                {{ item['text'] | trim }}
            {%- else -%}
                {{ raise_exception("Unknown content type: " + item['type']) }}
            {%- endif -%}
        {%- endfor -%}
    {%- endif -%}
    
    {#- Close tool response wrapper -#}
    {%- if message['role'] == 'tool' -%}
        {{ '\n</tool_response>' -}}
    {%- endif -%}
    
    {#- End message turn -#}
    {{ '<end_of_turn>\n' -}}
{%- endfor -%}

{#- Add generation prompt if requested -#}
{%- if add_generation_prompt -%}
    {{ '<start_of_turn>model\n' }}
{%- endif -%}
Google org

Hi @tarruda ,

This looks incredibly useful, especially how you have adapted it to handle the official JSON- style function calling as described in the Google AI docs. The extra details , like validating the conversation roles and handling the tool response wrappers, make it very robust.

This will surely be a huge help of others who are looking to implement tool use with Gemma models . Great work. Thank you.

@lkv Most of the work has been done by the vLLM team.

In any case, I appreciate the compliments. Thanks!

Sign up or log in to comment