[BFCL] Question regarding Custom Handler implementation for Qwen2.5-Instruct fine-tuned on Rlla-4k (BFCL-v4)

import json
import re
from typing import Any

from bfcl_eval.model_handler.local_inference.base_oss_handler import OSSHandler
from bfcl_eval.model_handler.utils import convert_to_function_call
from overrides import override

class Qwen2RLLAHandler(OSSHandler):
    def __init__(
        self,
        model_name,
        temperature,
        registry_name,
        is_fc_model,
        dtype="bfloat16",
        **kwargs,
    ) -> None:
        super().__init__(model_name, temperature, registry_name, is_fc_model, **kwargs)
        self.model_name_huggingface = model_name

    @override
    def decode_ast(self, result, language, has_tool_call_tag):
        tool_calls = self._extract_tool_calls(result)
        if type(tool_calls) != list or any(type(item) != dict for item in tool_calls):
            raise ValueError(f"Model did not return a list of function calls: {result}")
        return [
            {call["name"]: {k: v for k, v in call["arguments"].items()}}
            for call in tool_calls
        ]

    @override
    def decode_execute(self, result, has_tool_call_tag):
        tool_calls = self._extract_tool_calls(result)
        if type(tool_calls) != list or any(type(item) != dict for item in tool_calls):
            raise ValueError(f"Model did not return a list of function calls: {result}")
        decoded_result = []
        for item in tool_calls:
            if type(item) == str:
                item = eval(item)
            decoded_result.append({item["name"]: item["arguments"]})
        return convert_to_function_call(decoded_result)

    @override
    def _format_prompt(self, messages, function):
        formatted_prompt = ""


        if len(function) > 0:
            system_content = "You are a helpful multi-turn dialogue assistant capable of leveraging tool calls to solve user tasks and provide structured chat responses.\n\n**Available Tools**\nIn your response, you can use the following tools:\n"
            

            for idx, tool in enumerate(function, 1):
                name = tool.get("name", "")
                desc = tool.get("description", "")

                params_str = json.dumps(tool.get("parameters", {}), ensure_ascii=False)
                system_content += f"{idx}. Name: {name}\nDescription: {desc}\nParameters: {params_str}\n"

            system_content += "\n**Steps for Each Turn**\n1. **Think:** Recall relevant context and analyze the current user goal.\n2. **Decide on Tool Usage:** If a tool is needed, specify the tool and its parameters.\n3. **Respond Appropriately:** If a response is needed, generate one while maintaining consistency across user queries.\n\n**Output Format**\n"
            
            system_content += "```" + "plaintext\n<think> Your thoughts and reasoning </think>\n<tool_call>\n"
            system_content += "{\"name\": \"Tool name\", \"parameters\": {\"Parameter name\": \"Parameter content\", \"... ...\": \"... ...\"}}\n"
            system_content += "{\"name\": \"... ...\", \"parameters\": {\"... ...\": \"... ...\", \"... ...\": \"... ...\"}}\n...\n</tool_call>\n"
            system_content += "<response> AI's final response </response>\n" + "```" + "\n\n"
            
            system_content += "**Important Notes**\n1. You must always include the `<think>` field to outline your reasoning. Provide at least one of `<tool_call>` or `<response>`. Decide whether to use `<tool_call>` (possibly multiple times), `<response>`, or both.\n2. You can invoke multiple tool calls simultaneously in the `<tool_call>` fields. Each tool call should be a JSON object with a \"name\" field and an \"parameters\" field containing a dictionary of parameters. If no parameters are needed, leave the \"parameters\" field an empty dictionary.\n3. Refer to the previous dialogue records in the history, including the user's queries, previous `<tool_call>`, `<response>`, and any tool feedback noted as `<obs>` (if exists)."
            
            formatted_prompt += f"<|im_start|>system\n{system_content}<|im_end|>\n"
        else:
            if messages[0]["role"] == "system":
                formatted_prompt += f"<|im_start|>system\n{messages[0]['content']}<|im_end|>\n"

        is_first_user = True
        
        for idx, message in enumerate(messages):
            role = message["role"]
            content = message.get("content", "")

            if role == "system":
                continue

            elif role == "user":
                formatted_prompt += f"<|im_start|>user\n"
                if is_first_user:
                    formatted_prompt += f"**Dialogue Records History**\n"
                    is_first_user = False
                formatted_prompt += f"<user> {content} </user>\n<|im_end|>\n"

            elif role == "assistant":
                formatted_prompt += f"<|im_start|>assistant\n"

                if "reasoning_content" in message and message["reasoning_content"]:
                    formatted_prompt += f"<think> {message['reasoning_content']} </think>\n"

                if content:
                    formatted_prompt += f"<response> {content} </response>\n"
                
                if "tool_calls" in message:
                    formatted_prompt += "<tool_call>\n"
                    for tool_call in message["tool_calls"]:
                        if "function" in tool_call:
                            tool_call = tool_call["function"]
                        
                        args = tool_call.get("arguments", {})
                        if isinstance(args, str):
                            try:
                                args = json.loads(args)
                            except:
                                args = {}
                                
                        tool_str = json.dumps({"name": tool_call["name"], "parameters": args}, ensure_ascii=False)
                        formatted_prompt += f"{tool_str}\n"
                    formatted_prompt += "</tool_call>\n"

                formatted_prompt += "<|im_end|>\n"

            elif role == "tool":

                formatted_prompt += f"<|im_start|>user\n<obs> {content} </obs>\n<|im_end|>\n"

        formatted_prompt += "<|im_start|>assistant\n"
        
        return formatted_prompt

    @override
    def _pre_query_processing_prompting(self, test_entry: dict) -> dict:
        functions: list = test_entry["function"]
        return {"message": [], "function": functions}

    @override
    def _parse_query_response_prompting(self, api_response: Any) -> dict:
        model_response = api_response.choices[0].text
        
        reasoning_content = ""
        cleaned_response = model_response
        if "</think>" in model_response:
            parts = model_response.split("</think>")
            reasoning_content = parts[0].rstrip("\n").split("<think>")[-1].lstrip("\n")
            cleaned_response = parts[-1].lstrip("\n")
            
        extracted_tool_calls = self._extract_tool_calls(cleaned_response)

        if len(extracted_tool_calls) > 0:
            model_responses_message_for_chat_history = {
                "role": "assistant",
                "content": "",
                "tool_calls": extracted_tool_calls,
                "reasoning_content": reasoning_content,
            }
        else:
            model_responses_message_for_chat_history = {
                "role": "assistant",
                "content": cleaned_response.strip(),
                "reasoning_content": reasoning_content,
            }

        return {
            "model_responses": cleaned_response.strip(),
            "reasoning_content": reasoning_content,
            "model_responses_message_for_chat_history": model_responses_message_for_chat_history,
            "input_token": api_response.usage.prompt_tokens,
            "output_token": api_response.usage.completion_tokens,
        }

    @override
    def _add_assistant_message_prompting(
        self, inference_data: dict, model_response_data: dict
    ) -> dict:
        inference_data["message"].append(
            model_response_data["model_responses_message_for_chat_history"],
        )
        return inference_data

    @staticmethod
    def _extract_tool_calls(input_string):
        pattern = r"<tool_call>\n?(.*?)\n?</tool_call>"
        matches = re.findall(pattern, input_string, re.DOTALL)

        result = []
        for match in matches:

            lines = match.strip().split('\n')
            for line in lines:
                line = line.strip()
                if not line:
                    continue
                try:
                    tool_obj = json.loads(line)
                    arguments = tool_obj.get("parameters", {})
                    if isinstance(arguments, str):
                        arguments = json.loads(arguments)
                        
                    result.append({
                        "name": tool_obj.get("name", ""),
                        "arguments": arguments
                    })
                except Exception as e:
                    pass
        return result