NexusRaven-13B, a new SOTA Open-Source LLM for function calling. This repo contains everything for reproducing our evaluation on NexusRaven-13B and baselines.
from os import environ
from os.path import abspath, dirname, join


ROOT_DIR = dirname(abspath(__file__))
CACHE_DIR = abspath(join(ROOT_DIR, "..", "cache"))

environ["HF_DATASETS_CACHE"] = CACHE_DIR

raven/__init__.py
nexusflowai/NexusRaven
import os

import json

import re

from collections import defaultdict

import requests

from datasets import Dataset

from raven import ROOT_DIR


class ToolAlpacaDataPostprocessHelper:
    def run(self) -> None:
        response = requests.get(
            "https://raw.githubusercontent.com/tangqiaoyu/ToolAlpaca/main/data/eval_simulated.json"
        )
        content = response.content

        file = json.loads(content)
        eval_dataset = defaultdict(list)
        for item in file:
            (
                docstring,
                function_mapping,
            ) = self.generate_valid_python_functions_with_types(
                item["Function_Description"]
            )
            for instruction, golden_turns in zip(
                item["Instructions"], item["Golden_Answers"]
            ):
                # Transform the golden answers into a function call
                # We only want single turn for now.
                include = True
                # Write to the file
                parsed_turns = []
                for golden_answer in golden_turns:
                    golden_name = golden_answer["Action"]

                    if not golden_name in item["Function_Description"]:
                        print("ERROR", golden_name)
                        include = False
                        continue

                    golden_name = golden_name.replace("-", "_")
                    golden_call = golden_answer["Action_Input"]

                    golden_call = golden_call.replace(
                        """"symbols": "symbols": """, """"symbols": """
                    )
                    function_call = function_mapping[golden_name]
                    try:
                        function_call = self.construct_python_call(
                            function_call, golden_call
                        )
                    except json.decoder.JSONDecodeError:
                        # Skip this sample
                        continue
                    parsed_turns.append(function_call)

                # Serialize the list into a JSON string
                if len(parsed_turns) == 0:
                    continue

                serialized_parsed_turns = json.dumps(parsed_turns)
                serialized_backup = json.dumps((instruction, golden_turns))
                call = {
                    "docstring": docstring,
                    "instruction": instruction,
                    "golden_answer": serialized_parsed_turns,
                    "backup": serialized_backup,
                }
                key = f"turns_{len(golden_turns)}"
                if include:
                    eval_dataset[key].append(call)

        # Check the length of each list in the dict
        lengths = [len(v) for v in eval_dataset.values()]

        # If they aren't all the same, find the max length and pad the shorter ones
        if len(set(lengths)) > 1:
            max_length = max(lengths)
            for key, value in eval_dataset.items():
                if len(value) < max_length:
                    eval_dataset[key] = value + [{}] * (max_length - len(value))

        # Now the lengths should be the same and Dataset.from_dict() should work
        hf_dataset = Dataset.from_dict(eval_dataset)

        output_path = os.path.join(
            ROOT_DIR, "data", "resources", "toolalpaca_queries.json"
        )
        with open(output_path, "w") as f:
            json.dump(list(hf_dataset.to_iterable_dataset()), f, indent=4)

    def construct_python_call(self, function_call: str, golden_call: str) -> str:
        # Step 1: Parsing function_call
        function_match = re.match(r"(\w+)\((.*)\)", function_call)
        if not function_match:
            raise ValueError("Invalid function_call format")

        function_name = function_match.group(1)
        function_args_str = function_match.group(2).strip()

        # Handle no-argument case
        if not function_args_str:
            return f"{function_name}()"

        # Extract argument names
        arg_pairs = function_args_str.split(",")
        arg_names = [arg_pair.split(":")[0].strip() for arg_pair in arg_pairs]

        # Step 2: Parsing golden_call
        arg_values_dict = json.loads(golden_call)

        # Step 3: Generating actual Python call
        actual_call_args = []
        for arg_name in arg_names:
            # Default to None if the argument name is not found in arg_values_dict
            value = arg_values_dict.get(arg_name, None)

            if isinstance(value, str):
                actual_call_args.append(f'{arg_name}="{value}"')
            else:
                actual_call_args.append(f"{arg_name}={value}")

        actual_call = f'{function_name}({", ".join(actual_call_args)})'
        return actual_call

    def generate_valid_python_functions_with_types(self, functions):
        """
        Modify the function to use valid Python types for the parameters
        """
        python_code = ""
        type_mapping = {
            "Array[Object": "list[dict]",
            "Array": "list",
            "string": "str",
            "integer": "int",
            "array": "list",
            "object": "dict",
            "boolean": "bool",
            "Object.": "dict",
            "number": "float",
        }
        function_mapping = {}

        for function_name, function_desc in functions.items():
            if function_desc == "":
                continue

            function_name = function_name.replace("-", "_")

            summary = function_desc.split("\n")[0]

            operation_id = function_name
            try:
                parameters = json.loads(
                    function_desc.split("\n")[1].replace("Parameters: ", "")
                )
            except:  # pylint: disable=bare-except
                # Try using normal
                try:
                    parameters = json.loads(
                        function_desc.split("\n")[0].replace("Parameters: ", "")
                    )
                except:  # pylint: disable=bare-except
                    continue
            full_parameters = []

            function_name = operation_id
            function_parameters = []
            for param_name, param in parameters.items():
                param_type = None
                for key, value in type_mapping.items():
                    if key in param:
                        if param_type:
                            # Multiple types
                            raise ValueError
                        param_type = value
                        break
                if not param_type:
                    # No types
                    print(param)
                    raise ValueError

                param_default = None
                match = re.search(r"default is (.+)", param)
                if match:
                    extracted_value = match.group(1)
                    param_default = extracted_value
                param_required = (
                    not "Optional." in param
                )  # "required" in param or "Required" in param

                if param_required:
                    function_parameters.append(f"{param_name}: {param_type}")
                else:
                    function_parameters.append(
                        f"{param_name}: {param_type} = {param_default}"
                    )

                full_parameters.append(
                    {
                        "name": param_name,
                        "description": param,
                        "param_type": param_type,
                        "required": param_required,
                        "default": param_default,
                    }
                )

            function_parameters_str = ", ".join(function_parameters)

            docstring_lines = [
                '"""',
                f"{summary}",
                "",
            ]
            if len(parameters) > 0:
                docstring_lines.append("Parameters:")
            for param in full_parameters:
                param_name = param["name"]
                param_desc = param.get("description", "")
                param_type = param["param_type"]
                param_required = (
                    " required" if param.get("required", False) else " optional"
                )
                docstring_lines.append(
                    f"    {param_name} ({param_type}){param_required}: {param_desc}"
                )

            docstring_lines.append('"""')
            docstring_str = "\n".join(docstring_lines)

            if function_name == "":
                continue

            function_prototype = f"\nOPTION:\n<func_start>def {function_name}({function_parameters_str})<func_end>\n<docstring_start>\n{docstring_str}\n<docstring_end>\n"
            function_mapping[
                function_name
            ] = f"{function_name}({function_parameters_str})"

            python_code += function_prototype

        return python_code, function_mapping


if __name__ == "__main__":
    ToolAlpacaDataPostprocessHelper().run()

raven/data/process_toolalpaca_evaluation_data.py
nexusflowai/NexusRaven
def searchCVE(
    cpeName=None,
    cveId=None,
    cvssV2Metrics=None,
    cvssV2Severity=None,
    cvssV3Metrics=None,
    cvssV3Severity=None,
    cweId=None,
    hasCertAlerts=None,
    hasCertNotes=None,
    hasKev=None,
    hasOval=None,
    isVulnerable=None,
    keywordExactMatch=None,
    keywordSearch=None,
    lastModStartDate=None,
    lastModEndDate=None,
    noRejected=None,
    pubStartDate=None,
    pubEndDate=None,
    sourceIdentifier=None,
    versionEnd=None,
    versionEndType=None,
    versionStart=None,
    versionStartType=None,
    virtualMatchString=None,
    limit=None,
    delay=None,
    key=None,
    verbose=None,
):
    """
    Build and send GET request then return list of objects containing a collection of CVEs. For more information on the parameters available, please visit https://nvd.nist.gov/developers/vulnerabilities

    Args:
    cpeName (str): Please do not confuse this with keywordSearch; this requires the argument to start with "cpe", whereas the keywordSearch argument allows for arbitrary keywords. This value will be compared agains the CPE Match Criteria within a CVE applicability statement. (i.e. find the vulnerabilities attached to that CPE). Partial match strings are allowed.

    cveId (str): Please pass in a string integer, like "1" or "30". Returns a single CVE that already exists in the NVD.

    cvssV2Metrics (str): This parameter returns only the CVEs that match the provided CVSSv2 vector string. Either full or partial vector strings may be used. This parameter cannot be used in requests that include cvssV3Metrics.

    cvssV2Severity (str): Find vulnerabilities having a LOW, MEDIUM, or HIGH version 2 severity.

    cvssV3Metrics (str): This parameter returns only the CVEs that match the provided CVSSv3 vector string. Either full or partial vector strings may be used. This parameter cannot be used in requests that include cvssV2Metrics.

    cvssV3Severity (str): Find vulnerabilities having a LOW, MEDIUM, HIGH, or CRITICAL version 3 severity.

    cweId (str): Please pass in a string integer, like "1" or "30". Filter collection by CWE (Common Weakness Enumeration) ID. You can find a list at https://cwe.mitre.org/. A CVE can have multiple CWE IDs assigned to it.

    hasCertAlerts (bool): Returns CVE that contain a Technical Alert from US-CERT.

    hasCertNotes (bool): Returns CVE that contain a Vulnerability Note from CERT/CC.

    hasOval (bool): Returns CVE that contain information from MITRE's Open Vulnerability and Assessment Language (OVAL) before this transitioned to the Center for Internet Security (CIS).

    isVulnerable (bool): Returns CVE associated with a specific CPE, where the CPE is also considered vulnerable. REQUIRES cpeName parameter. isVulnerable is not compatible with virtualMatchString parameter.

    keywordExactMatch (bool): When keywordSearch is used along with keywordExactmatch, it will search the NVD for CVEs containing exactly what was passed to keywordSearch. REQUIRES keywordSearch.

    keywordSearch (str): Searches CVEs where a word or phrase is found in the current description. If passing multiple keywords with a space character in between then each word must exist somewhere in the description, not necessarily together unless keywordExactMatch=True is passed to searchCVE.

    lastModStartDate (str, datetime obj): These parameters return only the CVEs that were last modified during the specified period. If a CVE has been modified more recently than the specified period, it will not be included in the response. If filtering by the last modified date, both lastModStartDate and lastModEndDate are REQUIRED. The maximum allowable range when using any date range parameters is 120 consecutive days.

    lastModEndDate (str, datetime obj): Required if using lastModStartDate.

    noRejected (bool): Filters out all CVEs that are in a reject or rejected status. Searches without this parameter include rejected CVEs.

    pubStartDate (str, datetime obj): These parameters return only the CVEs that were added to the NVD (i.e., published) during the specified period. If filtering by the published date, both pubStartDate and pubEndDate are REQUIRED. The maximum allowable range when using any date range parameters is 120 consecutive days.

    pubEndDate (str, datetime obj): Required if using pubStartDate.

    sourceIdentifier (str): Returns CVE where the data source of the CVE is the value that is passed to sourceIdentifier.

    versionEnd (str): Must be combined with versionEndType and virtualMatchString. Returns only the CVEs associated with CPEs in specific version ranges.

    versionEndType (str): Must be combined with versionEnd and virtualMatchString. Valid values are including or excluding. Denotes to include the specified version in versionEnd, or exclude it.

    versionStart (str): Must be combined with versionStartType and virtualMatchString. Returns only CVEs with specific versions. Requests that include versionStart cannot include a version component in the virtualMatchString.

    versionStartType (str): Must be combined with versionStart and virtualMatchString. Valid values are including or excluding. Denotes to include the specified version in versionStart, or exclude it.

    virtualMatchString (str): A more broad filter compared to cpeName. The cpe match string that is passed to virtualMatchString is compared against the CPE Match Criteria present on CVE applicability statements.

    limit (int): Custom argument to limit the number of results of the search. Allowed any number between 1 and 2000.

    delay (int): Can only be used if an API key is provided. This allows the user to define a delay. The delay must be greater than 0.6 seconds. The NVD API recommends scripts sleep for atleast 6 seconds in between requests.

    key (str): NVD API Key. Allows for the user to define a delay. NVD recommends scripts sleep 6 seconds in between requests. If no valid API key is provided, requests are sent with a 6 second delay.

    verbose (bool): Prints the URL request for debugging purposes.
    """


def searchCPE(
    cpeNameId=None,
    cpeMatchString=None,
    keywordExactMatch=None,
    keywordSearch=None,
    lastModStartDate=None,
    lastModEndDate=None,
    matchCriteriaId=None,
    limit=None,
    key=None,
    delay=None,
    verbose=None,
):
    """
    Build and send GET request then return list of objects containing a collection of CPEs.

    Args:
    cpeNameId (str): Returns a specific CPE record using its UUID. If a correctly formatted UUID is passed but it does not exist, it will return empty results. The UUID is the cpeNameId value when searching CPE.

    cpeMatchString (str): Use a partial CPE name to search for other CPE names.

    keywordExactMatch (bool): Searches metadata within CPE title and reference links for an exact match of the phrase or word passed to it. Must be included with keywordSearch.

    keywordSearch (str): Returns CPE records where a word or phrase is found in the metadata title or reference links. Space characters act as an AND statement.

    lastModStartDate (str/datetime obj): CPE last modification start date. Maximum 120 day range. A start and end date is required. All times are in UTC 00:00. A datetime object or string can be passed as a date. NVDLib will automatically parse the datetime object into the correct format. String Example: ‘2020-06-28 00:00’

    lastModEndDate (str/datetime obj): CPE last modification end date. Maximum 120 day range. Must be included with lastModStartDate. Example: ‘2020-06-28 00:00’

    limit (int): Limits the number of results of the search.

    key (str): NVD API Key. Allows for a request every 0.6 seconds instead of 6 seconds.

    delay (int): Can only be used if an API key is provided. The amount of time to sleep in between requests. Must be a value above 0.6 seconds if an API key is present. delay is set to 6 seconds if no API key is passed.

    verbose (bool): Prints the URL request for debugging purposes.
    """

raven/data/resources/cve_cpe_function_definitions.py
nexusflowai/NexusRaven
# pylint: disable-all
"""
This is an exact copy of https://github.com/OpenBMB/ToolBench/blob/master/toolbench/inference/Downstream_tasks/rapidapi.py
except the output results are written to `results.txt`. (see line 688)
"""
import re
import os
import json
import time
import requests
from tqdm import tqdm
from termcolor import colored
import random
from toolbench.inference.LLM.chatgpt_function_model import ChatGPTFunction
from toolbench.inference.LLM.davinci_model import Davinci
from toolbench.inference.LLM.tool_llama_lora_model import ToolLLaMALoRA
from toolbench.inference.LLM.tool_llama_model import ToolLLaMA
from toolbench.inference.LLM.retriever import ToolRetriever
from toolbench.inference.Algorithms.single_chain import single_chain
from toolbench.inference.Algorithms.DFS import DFS_tree_search
from toolbench.inference.server import get_rapidapi_response
from toolbench.utils import standardize, change_name, replace_llama_with_condense

from toolbench.inference.Downstream_tasks.base_env import base_env


# For pipeline environment preparation
def get_white_list(tool_root_dir):
    # print(tool_root_dir)
    white_list_dir = os.path.join(tool_root_dir)
    white_list = {}
    for cate in tqdm(os.listdir(white_list_dir)):
        if not os.path.isdir(os.path.join(white_list_dir, cate)):
            continue
        for file in os.listdir(os.path.join(white_list_dir, cate)):
            if not file.endswith(".json"):
                continue
            standard_tool_name = file.split(".")[0]
            # print(standard_tool_name)
            with open(os.path.join(white_list_dir, cate, file)) as reader:
                js_data = json.load(reader)
            origin_tool_name = js_data["tool_name"]
            white_list[standardize(origin_tool_name)] = {
                "description": js_data["tool_description"],
                "standard_tool_name": standard_tool_name,
            }
    return white_list


def contain(candidate_list, white_list):
    output = []
    for cand in candidate_list:
        if cand not in white_list.keys():
            return False
        output.append(white_list[cand])
    return output


# rapidapi env wrapper
class rapidapi_wrapper(base_env):
    def __init__(self, query_json, tool_descriptions, retriever, args, process_id=0):
        super(rapidapi_wrapper).__init__()

        self.tool_root_dir = args.tool_root_dir
        self.toolbench_key = args.toolbench_key
        self.rapidapi_key = args.rapidapi_key
        self.use_rapidapi_key = args.use_rapidapi_key
        self.api_customization = args.api_customization
        self.service_url = "http://8.218.239.54:8080/rapidapi"
        self.max_observation_length = args.max_observation_length
        self.observ_compress_method = args.observ_compress_method
        self.retriever = retriever
        self.process_id = process_id

        self.tool_names = []
        self.cate_names = []

        self.input_description = query_json["query"]
        self.functions = []
        self.api_name_reflect = {}

        if self.retriever is not None:
            query_json = self.retrieve_rapidapi_tools(
                self.input_description, args.retrieved_api_nums, args.tool_root_dir
            )
            data_dict = self.fetch_api_json(query_json)
            tool_descriptions = self.build_tool_description(data_dict)
        else:
            data_dict = self.fetch_api_json(query_json)

        for k, api_json in enumerate(data_dict["api_list"]):
            standard_tool_name = tool_descriptions[k][0]
            (
                openai_function_json,
                cate_name,
                pure_api_name,
            ) = self.api_json_to_openai_json(api_json, standard_tool_name)
            self.functions.append(openai_function_json)

            self.api_name_reflect[openai_function_json["name"]] = pure_api_name
            self.tool_names.append(standard_tool_name)
            self.cate_names.append(cate_name)

        finish_func = {
            "name": "Finish",
            "description": "If you believe that you have obtained a result that can answer the task, please call this function to provide the final answer. Alternatively, if you recognize that you are unable to proceed with the task in the current state, call this function to restart. Remember: you must ALWAYS call this function at the end of your attempt, and the only part that will be shown to the user is the final answer, so it should contain sufficient information.",
            "parameters": {
                "type": "object",
                "properties": {
                    "return_type": {
                        "type": "string",
                        "enum": ["give_answer", "give_up_and_restart"],
                    },
                    "final_answer": {
                        "type": "string",
                        "description": 'The final answer you want to give the user. You should have this field if "return_type"=="give_answer"',
                    },
                },
                "required": ["return_type"],
            },
        }

        self.functions.append(finish_func)
        self.CALL_MAX_TIME = 3
        self.task_description = f"""You should use functions to help handle the real time user querys. Remember:
1.ALWAYS call \"Finish\" function at the end of the task. And the final answer should contain enough information to show to the user,If you can't handle the task, or you find that function calls always fail(the function is not valid now), use function Finish->give_up_and_restart.
2.Do not use origin tool names, use only subfunctions' names.
You have access of the following tools:\n"""

        unduplicated_reflection = {}
        for standardize_tool_name, tool_des in tool_descriptions:
            unduplicated_reflection[standardize_tool_name] = tool_des

        for k, (standardize_tool_name, tool_des) in enumerate(
            unduplicated_reflection.items()
        ):
            striped = tool_des[:512].replace("\n", "").strip()
            if striped == "":
                striped = "None"
            self.task_description += f"{k+1}.{standardize_tool_name}: {striped}\n"

        self.success = 0

    def build_tool_description(self, data_dict):
        white_list = get_white_list(self.tool_root_dir)
        origin_tool_names = [
            standardize(cont["tool_name"]) for cont in data_dict["api_list"]
        ]
        tool_des = contain(origin_tool_names, white_list)
        tool_descriptions = [
            [cont["standard_tool_name"], cont["description"]] for cont in tool_des
        ]
        return tool_descriptions

    def retrieve_rapidapi_tools(self, query, top_k, jsons_path):
        retrieved_tools = self.retriever.retrieving(query, top_k=top_k)
        query_json = {"api_list": []}
        for tool_dict in retrieved_tools:
            if len(query_json["api_list"]) == top_k:
                break
            category = tool_dict["category"]
            tool_name = tool_dict["tool_name"]
            api_name = tool_dict["api_name"]
            if os.path.exists(jsons_path):
                if os.path.exists(os.path.join(jsons_path, category)):
                    if os.path.exists(
                        os.path.join(jsons_path, category, tool_name + ".json")
                    ):
                        query_json["api_list"].append(
                            {
                                "category_name": category,
                                "tool_name": tool_name,
                                "api_name": api_name,
                            }
                        )
        return query_json

    def fetch_api_json(self, query_json):
        data_dict = {"api_list": []}
        for item in query_json["api_list"]:
            cate_name = item["category_name"]
            tool_name = standardize(item["tool_name"])
            api_name = change_name(standardize(item["api_name"]))
            tool_json = json.load(
                open(
                    os.path.join(self.tool_root_dir, cate_name, tool_name + ".json"),
                    "r",
                )
            )
            append_flag = False
            api_dict_names = []
            for api_dict in tool_json["api_list"]:
                api_dict_names.append(api_dict["name"])
                pure_api_name = change_name(standardize(api_dict["name"]))
                if pure_api_name != api_name:
                    continue
                api_json = {}
                api_json["category_name"] = cate_name
                api_json["api_name"] = api_dict["name"]
                api_json["api_description"] = api_dict["description"]
                api_json["required_parameters"] = api_dict["required_parameters"]
                api_json["optional_parameters"] = api_dict["optional_parameters"]
                api_json["tool_name"] = tool_json["tool_name"]
                data_dict["api_list"].append(api_json)
                append_flag = True
                break
            if not append_flag:
                print(api_name, api_dict_names)
        return data_dict

    def api_json_to_openai_json(self, api_json, standard_tool_name):
        description_max_length = 256
        templete = {
            "name": "",
            "description": "",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": [],
                "optional": [],
            },
        }

        map_type = {"NUMBER": "integer", "STRING": "string", "BOOLEAN": "boolean"}

        pure_api_name = change_name(standardize(api_json["api_name"]))
        templete["name"] = pure_api_name + f"_for_{standard_tool_name}"
        templete["name"] = templete["name"][-64:]

        templete[
            "description"
        ] = f'This is the subfunction for tool "{standard_tool_name}", you can use this tool.'

        if api_json["api_description"].strip() != "":
            tuncated_description = (
                api_json["api_description"]
                .strip()
                .replace(api_json["api_name"], templete["name"])[
                    :description_max_length
                ]
            )
            templete["description"] = (
                templete["description"]
                + f'The description of this function is: "{tuncated_description}"'
            )
        if (
            "required_parameters" in api_json.keys()
            and len(api_json["required_parameters"]) > 0
        ):
            for para in api_json["required_parameters"]:
                name = standardize(para["name"])
                name = change_name(name)
                if para["type"] in map_type:
                    param_type = map_type[para["type"]]
                else:
                    param_type = "string"
                prompt = {
                    "type": param_type,
                    "description": para["description"][:description_max_length],
                }

                default_value = para["default"]
                if len(str(default_value)) != 0:
                    prompt = {
                        "type": param_type,
                        "description": para["description"][:description_max_length],
                        "example_value": default_value,
                    }
                else:
                    prompt = {
                        "type": param_type,
                        "description": para["description"][:description_max_length],
                    }

                templete["parameters"]["properties"][name] = prompt
                templete["parameters"]["required"].append(name)
            for para in api_json["optional_parameters"]:
                name = standardize(para["name"])
                name = change_name(name)
                if para["type"] in map_type:
                    param_type = map_type[para["type"]]
                else:
                    param_type = "string"

                default_value = para["default"]
                if len(str(default_value)) != 0:
                    prompt = {
                        "type": param_type,
                        "description": para["description"][:description_max_length],
                        "example_value": default_value,
                    }
                else:
                    prompt = {
                        "type": param_type,
                        "description": para["description"][:description_max_length],
                    }

                templete["parameters"]["properties"][name] = prompt
                templete["parameters"]["optional"].append(name)

        return templete, api_json["category_name"], pure_api_name

    def check_success(self):
        return self.success

    def to_json(self):
        return {}

    def restart(self):
        pass

    def get_score(self):
        return 0.0

    def step(self, **args):
        obs, code = self._step(**args)
        if len(obs) > self.max_observation_length:
            obs = obs[: self.max_observation_length] + "..."
        return obs, code

    def _step(self, action_name="", action_input=""):
        """Need to return an observation string and status code:
        0 means normal response
        1 means there is no corresponding api name
        2 means there is an error in the input
        3 represents the end of the generation and the final answer appears
        4 means that the model decides to pruning by itself
        5 represents api call timeout
        6 for 404
        7 means not subscribed
        8 represents unauthorized
        9 represents too many requests
        10 stands for rate limit
        11 message contains "error" field
        12 error sending request
        """
        if action_name == "Finish":
            try:
                json_data = json.loads(action_input, strict=False)
            except:
                json_data = {}
                if '"return_type": "' in action_input:
                    if '"return_type": "give_answer"' in action_input:
                        return_type = "give_answer"
                    elif '"return_type": "give_up_and_restart"' in action_input:
                        return_type = "give_up_and_restart"
                    else:
                        return_type = action_input[
                            action_input.find('"return_type": "')
                            + len('"return_type": "') : action_input.find('",')
                        ]
                    json_data["return_type"] = return_type
                if '"final_answer": "' in action_input:
                    final_answer = action_input[
                        action_input.find('"final_answer": "')
                        + len('"final_answer": "') :
                    ]
                    json_data["final_answer"] = final_answer
            if "return_type" not in json_data.keys():
                return '{error:"must have "return_type""}', 2
            if json_data["return_type"] == "give_up_and_restart":
                return '{"response":"chose to give up and restart"}', 4
            elif json_data["return_type"] == "give_answer":
                if "final_answer" not in json_data.keys():
                    return '{error:"must have "final_answer""}', 2

                self.success = 1  # succesfully return final_answer
                return '{"response":"successfully giving the final answer."}', 3
            else:
                return '{error:""return_type" is not a valid choice"}', 2
        else:
            for k, function in enumerate(self.functions):
                if function["name"].endswith(action_name):
                    pure_api_name = self.api_name_reflect[function["name"]]
                    payload = {
                        "category": self.cate_names[k],
                        "tool_name": self.tool_names[k],
                        "api_name": pure_api_name,
                        "tool_input": action_input,
                        "strip": self.observ_compress_method,
                        "toolbench_key": self.toolbench_key,
                    }
                    if self.process_id == 0:
                        print(
                            colored(
                                f"query to {self.cate_names[k]}-->{self.tool_names[k]}-->{action_name}",
                                color="yellow",
                            )
                        )
                    if self.use_rapidapi_key or self.api_customization:
                        payload["rapidapi_key"] = self.rapidapi_key
                        response = get_rapidapi_response(
                            payload, api_customization=self.api_customization
                        )
                    else:
                        time.sleep(2)  # rate limit: 30 per minute
                        headers = {"toolbench_key": self.toolbench_key}
                        response = requests.post(
                            self.service_url, json=payload, headers=headers, timeout=15
                        )
                        if response.status_code != 200:
                            return (
                                json.dumps(
                                    {
                                        "error": f"request invalid, data error. status_code={response.status_code}",
                                        "response": "",
                                    }
                                ),
                                12,
                            )
                        try:
                            response = response.json()
                        except:
                            print(response)
                            return (
                                json.dumps(
                                    {
                                        "error": f"request invalid, data error",
                                        "response": "",
                                    }
                                ),
                                12,
                            )
                    # 1 Hallucinating function names
                    # 4 means that the model decides to pruning by itself
                    # 5 represents api call timeout
                    # 6 for 404
                    # 7 means not subscribed
                    # 8 represents unauthorized
                    # 9 represents too many requests
                    # 10 stands for rate limit
                    # 11 message contains "error" field
                    # 12 error sending request
                    if response["error"] == "API not working error...":
                        status_code = 6
                    elif response["error"] == "Unauthorized error...":
                        status_code = 7
                    elif response["error"] == "Unsubscribed error...":
                        status_code = 8
                    elif response["error"] == "Too many requests error...":
                        status_code = 9
                    elif response["error"] == "Rate limit per minute error...":
                        print("Reach api calling limit per minute, sleeping...")
                        time.sleep(10)
                        status_code = 10
                    elif response["error"] == "Message error...":
                        status_code = 11
                    else:
                        status_code = 0
                    return json.dumps(response), status_code
                    # except Exception as e:
                    #     return json.dumps({"error": f"Timeout error...{e}", "response": ""}), 5
            return (
                json.dumps(
                    {"error": f"No such function name: {action_name}", "response": ""}
                ),
                1,
            )


class pipeline_runner:
    def __init__(self, args, add_retrieval=False, process_id=0, server=False):
        self.args = args
        self.add_retrieval = add_retrieval
        self.process_id = process_id
        self.server = server
        if not self.server:
            self.task_list = self.generate_task_list()
        else:
            self.task_list = []

    def get_backbone_model(self):
        args = self.args
        if args.backbone_model == "toolllama":
            # ratio = 4 means the sequence length is expanded by 4, remember to change the model_max_length to 8192 (2048 * ratio) for ratio = 4
            replace_llama_with_condense(ratio=4)
            if args.lora:
                backbone_model = ToolLLaMALoRA(
                    base_name_or_path=args.model_path, model_name_or_path=args.lora_path
                )
            else:
                backbone_model = ToolLLaMA(model_name_or_path=args.model_path)
        else:
            backbone_model = args.backbone_model
        return backbone_model

    def get_retriever(self):
        return ToolRetriever(
            corpus_tsv_path=self.args.corpus_tsv_path,
            model_path=self.args.retrieval_model_path,
        )

    def get_args(self):
        return self.args

    def generate_task_list(self):
        args = self.args
        query_dir = args.input_query_file
        answer_dir = args.output_answer_file
        if not os.path.exists(answer_dir):
            os.mkdir(answer_dir)
        method = args.method
        backbone_model = self.get_backbone_model()
        white_list = get_white_list(args.tool_root_dir)
        task_list = []
        querys = json.load(open(query_dir, "r"))
        for query_id, data_dict in enumerate(querys):
            if "query_id" in data_dict:
                query_id = data_dict["query_id"]
            if "api_list" in data_dict:
                origin_tool_names = [
                    standardize(cont["tool_name"]) for cont in data_dict["api_list"]
                ]
                tool_des = contain(origin_tool_names, white_list)
                if tool_des == False:
                    continue
                tool_des = [
                    [cont["standard_tool_name"], cont["description"]]
                    for cont in tool_des
                ]
            else:
                tool_des = None
            task_list.append(
                (
                    method,
                    backbone_model,
                    query_id,
                    data_dict,
                    args,
                    answer_dir,
                    tool_des,
                )
            )
        return task_list

    def method_converter(
        self,
        backbone_model,
        openai_key,
        method,
        env,
        process_id,
        single_chain_max_step=12,
        max_query_count=60,
        callbacks=None,
    ):
        if callbacks is None:
            callbacks = []
        if backbone_model == "chatgpt_function":
            model = "gpt-3.5-turbo-16k-0613"
            llm_forward = ChatGPTFunction(model=model, openai_key=openai_key)
        elif backbone_model == "davinci":
            model = "text-davinci-003"
            llm_forward = Davinci(model=model, openai_key=openai_key)
        else:
            model = backbone_model
            llm_forward = model

        if method.startswith("CoT"):
            passat = int(method.split("@")[-1])
            chain = single_chain(llm=llm_forward, io_func=env, process_id=process_id)
            result = chain.start(
                pass_at=passat, single_chain_max_step=single_chain_max_step, answer=1
            )
        elif method.startswith("DFS"):
            pattern = r".+_w(\d+)"
            re_result = re.match(pattern, method)
            assert re_result != None
            width = int(re_result.group(1))
            with_filter = True
            if "woFilter" in method:
                with_filter = False
            chain = DFS_tree_search(
                llm=llm_forward, io_func=env, process_id=process_id, callbacks=callbacks
            )
            result = chain.start(
                single_chain_max_step=single_chain_max_step,
                tree_beam_size=width,
                max_query_count=max_query_count,
                answer=1,
                with_filter=with_filter,
            )
        else:
            print("invalid method")
            raise NotImplementedError
        return chain, result

    def run_single_task(
        self,
        method,
        backbone_model,
        query_id,
        data_dict,
        args,
        output_dir_path,
        tool_des,
        retriever=None,
        process_id=0,
        callbacks=None,
        server=None,
    ):
        if server is None:
            server = self.server
        if callbacks is None:
            if server:
                print("Warning: no callbacks are defined for server mode")
            callbacks = []
        splits = output_dir_path.split("/")
        os.makedirs("/".join(splits[:-1]), exist_ok=True)
        os.makedirs("/".join(splits), exist_ok=True)
        output_file_path = os.path.join(output_dir_path, f"{query_id}_{method}.json")
        if (not server) and os.path.exists(output_file_path):
            return
        [callback.on_tool_retrieval_start() for callback in callbacks]
        env = rapidapi_wrapper(
            data_dict, tool_des, retriever, args, process_id=process_id
        )
        [callback.on_tool_retrieval_end(tools=env.functions) for callback in callbacks]
        query = data_dict["query"]
        if process_id == 0:
            print(
                colored(
                    f"[process({process_id})]now playing {query}, with {len(env.functions)} APIs",
                    "green",
                )
            )
        [
            callback.on_request_start(
                user_input=query,
                method=method,
            )
            for callback in callbacks
        ]
        chain, result = self.method_converter(
            backbone_model=backbone_model,
            openai_key=args.openai_key,
            method=method,
            env=env,
            process_id=process_id,
            single_chain_max_step=12,
            max_query_count=200,
            callbacks=callbacks,
        )
        [
            callback.on_request_end(
                chain=chain.terminal_node[0].messages,
                outputs=chain.terminal_node[0].description,
            )
            for callback in callbacks
        ]
        if output_dir_path is not None:
            with open(output_file_path, "w") as writer:
                data = chain.to_json(answer=True, process=True)
                data["answer_generation"]["query"] = query
                json.dump(data, writer, indent=2)
                success = (
                    data["answer_generation"]["valid_data"]
                    and "give_answer" in data["answer_generation"]["final_answer"]
                )
                print(colored(f"[process({process_id})]valid={success}", "green"))
        return result

    def run(self):
        task_list = self.task_list
        random.seed(42)
        random.shuffle(task_list)
        print(f"total tasks: {len(task_list)}")
        new_task_list = []
        for task in task_list:
            out_dir_path = task[-2]
            query_id = task[2]
            output_file_path = os.path.join(
                out_dir_path, f"{query_id}_{self.args.method}.json"
            )
            if not os.path.exists(output_file_path):
                new_task_list.append(task)
        task_list = new_task_list
        print(f"undo tasks: {len(task_list)}")
        if self.add_retrieval:
            retriever = self.get_retriever()
        else:
            retriever = None
        for k, task in enumerate(task_list):
            print(
                f"process[{self.process_id}] doing task {k}/{len(task_list)}: real_task_id_{task[2]}"
            )

            """
            START Output query to `results.txt`
            """

            import json, time

            d_str = json.dumps(
                {
                    "task_id": task[2],
                    "query": task[3]["query"],
                    "timestamp": time.time(),
                }
            )
            with open("results.txt", "a+") as f:
                f.write(f"========\n{d_str}\n")

            """
            END Output query to `results.txt`
            """

            result = self.run_single_task(
                *task, retriever=retriever, process_id=self.process_id
            )

raven/data/resources/rapidapi.py
nexusflowai/NexusRaven
# pylint: disable-all
"""
This is an exact copy of https://github.com/OpenBMB/ToolBench/blob/master/toolbench/inference/server.py
except the results are output to `results.txt` (see line 109)
"""
from pydantic import BaseModel
import json
import os
from typing import Union
from toolbench.utils import standardize, change_name
import random


class Info(BaseModel):
    category: str
    tool_name: str
    api_name: str
    tool_input: Union[str, dict]
    strip: str


def prepare_tool_name_and_url(tools_root, info):
    category = info.category
    standard_category = category.replace(" ", "_").replace(",", "_").replace("/", "_")
    while " " in standard_category or "," in standard_category:
        standard_category = standard_category.replace(" ", "_").replace(",", "_")
    standard_category = standard_category.replace("__", "_")

    tool_name = info.tool_name
    api_name = change_name(standardize(info.api_name))
    if not tool_name.endswith(f"_for_{standard_category}"):
        tool_name = standardize(info.tool_name)
        code_string = f"""from {tools_root}.{standard_category}.{tool_name}.api import {api_name}"""
        tool_name += f"_for_{standard_category}"
    else:
        tmp_tool_name = standardize(tool_name.replace(f"_for_{standard_category}", ""))
        code_string = f"""from {tools_root}.{standard_category}.{tmp_tool_name}.api import {api_name}"""
    return tool_name, standard_category, api_name, code_string


def process_error(response):
    save_cache_flag = False
    switch_flag = False
    if (
        "The request to the API has timed out. Please try again later, or if the issue persists"
        in str(response)
    ):
        return_dict = {
            "error": "API temporarily not working error...",
            "response": response,
        }

    if "Your Client (working) ---> Gateway (working) ---> API (not working)" in str(
        response
    ):
        return_dict = {"error": "API not working error...", "response": response}

    elif "Unauthorized" in str(response) or "unauthorized" in str(response):
        save_cache_flag = True
        return_dict = {"error": "Unauthorized error...", "response": response}

    elif "You are not subscribed to this API." in str(response):
        switch_flag = True
        return_dict = {"error": "Unsubscribed error...", "response": response}

    elif "Too many requests" in str(response):
        switch_flag = True
        return_dict = {"error": "Too many requests error...", "response": response}

    elif "You have exceeded" in str(response) or "you are being rate limited" in str(
        response
    ):
        switch_flag = True
        return_dict = {"error": "Rate limit error...", "response": response}

    elif (
        "Access restricted. Check credits balance or enter the correct API key."
        in str(response)
    ):
        switch_flag = True
        return_dict = {"error": "Rate limit error...", "response": response}

    elif "Oops, an error in the gateway has occurred." in str(response):
        switch_flag = True
        return_dict = {"error": "Gateway error...", "response": response}

    elif "Blocked User. Please contact your API provider." in str(response):
        switch_flag = True
        return_dict = {"error": "Blocked error...", "response": response}

    elif "error" in str(response):
        return_dict = {"error": "Message error...", "response": response}

    else:
        save_cache_flag = True
        return_dict = {"error": "", "response": response}
    return return_dict, save_cache_flag, switch_flag


def run(toolbench_code_string, toolbench_api_name, toolbench_input_params_str):
    # get observation
    success_flag = False
    switch_flag = False
    save_cache = False
    exec(toolbench_code_string)
    try:
        eval_func_str = f"{toolbench_api_name}({toolbench_input_params_str})"

        """
        START Output query to to `results.txt`
        """

        import json, time

        d_str = json.dumps({"function_call": eval_func_str, "timestamp": time.time()})
        with open("results.txt", "a+") as f:
            f.write(f"{d_str}\n")

        """
        END Output query to to `results.txt`
        """

        new_func = eval(eval_func_str)
        response, save_cache, switch_flag = process_error(new_func)
        success_flag = True
    except Exception as e:
        response = {
            "error": f"Function executing {toolbench_code_string} error...\n{e}",
            "response": "",
        }
        save_cache = False
    return success_flag, switch_flag, response, save_cache


def dict_shorten(origin: dict, schema: dict):
    for key, value in list(origin.items()):
        if key not in schema:
            del origin[key]
        else:
            if isinstance(value, dict):
                dict_shorten(value, schema[key])  # schema[key] should be a dict
            elif isinstance(value, list):
                if value:
                    if isinstance(value[0], dict):
                        for item in value:
                            dict_shorten(
                                item, schema[key][0]
                            )  # schema[key] should be a list with only one dict element
    return origin


def observation_shorten(
    schema_root, response_dict, category, tool_name, api_name, strip_method
):
    print(random.random())
    if strip_method == "filter" or (strip_method == "random" and random.random() > 0.5):
        if isinstance(response_dict["response"], dict):
            if os.path.exists(os.path.join(schema_root, category)):
                if os.path.exists(
                    os.path.join(schema_root, category, tool_name + ".json")
                ):
                    schema_dicts = json.load(
                        open(
                            os.path.join(schema_root, category, tool_name + ".json"),
                            "r",
                        )
                    )
                    api_list = schema_dicts["api_list"]
                    schema = None
                    for schema_dict in api_list:
                        schema_api_name = change_name(standardize(schema_dict["name"]))
                        if (
                            schema_api_name == api_name
                            and len(schema_dict["schema"]) > 0
                        ):
                            schema = schema_dict["schema"]
                            break
                    if schema is not None:
                        response_dict["response"] = dict_shorten(
                            response_dict["response"], schema
                        )
    return str(response_dict["response"])


def get_rapidapi_response(
    input_dict: dict,
    api_customization: bool = False,
    tools_root: str = "data.toolenv.tools",
    schema_root: str = "data/toolenv/response_examples",
):
    info = Info
    info.category = input_dict["category"]
    info.tool_name = input_dict["tool_name"]
    info.api_name = input_dict["api_name"]
    info.tool_input = input_dict["tool_input"]
    info.strip = input_dict["strip"]
    rapidapi_key = input_dict["rapidapi_key"]

    tool_name, standard_category, api_name, code_string = prepare_tool_name_and_url(
        tools_root, info
    )
    tool_input = info.tool_input

    strip_method = info.strip

    try:
        tool_input = json.loads(tool_input)
    except Exception as e:
        if tool_input == "":
            tool_input = {}
        else:
            print(f"Can not parse tool input into json: {tool_input}")
            response_dict = {"error": f"Tool input parse error...\n", "response": ""}
            return response_dict

    input_params_str = ""
    if len(tool_input) > 0:
        for key, value in tool_input.items():
            if isinstance(value, str):
                input_params_str += f'{key}="{value}", '
            else:
                input_params_str += f"{key}={value}, "
    if not api_customization:
        input_params_str += f"toolbench_rapidapi_key='{rapidapi_key}'"
    success_flag, switch_flag, response_dict, save_cache = run(
        code_string, api_name, input_params_str
    )
    observation = observation_shorten(
        schema_root,
        response_dict,
        standard_category,
        tool_name.replace(f"_for_{standard_category}", ""),
        api_name,
        strip_method,
    )
    result = str(observation)[:2048]
    return {"error": response_dict["error"], "response": result}


if __name__ == "__main__":
    result = get_rapidapi_response(
        {
            "category": "Social",
            "tool_name": "olato_quotes",
            "api_name": "love_quote",
            "tool_input": "{}",
            "strip": "filter",
            "rapidapi_key": "",
        }
    )
    print(result)

raven/data/resources/server.py
nexusflowai/NexusRaven
def vt_get_dns_resolution_object(id: str, x_apikey: str):
    """
    This endpoint retrieves a Resolution object by its ID. A resolution object ID is made by appending the IP and the domain it resolves to together.

    Domain-IP resolutions. Resolution objects include the following attributes:
    date: <integer> date when the resolution was made (UTC timestamp).
    host_name: <string> domain or subdomain requested to the resolver.
    host_name_last_analysis_stats: <dictionary> last detection stats from the resolution's domain. Similar to the domains's last_analysis_stats attribute.
    ip_address: <string> IP address the domain was resolved to.
    ip_address_last_analysis_stats: <dictionary> last detection stats from the resolution's IP address. Similar to the IP address' last_analysis_stats attribute.
    resolver: <string> source of the resolution.

    Args:
    - id: string, required, Resolution object ID
    - x-apikey: string, required, Your API key
    """


def vt_get_objects_related_to_ip_address(
    ip: str, relationship: str, x_apikey: str, limit: int = None, cursor: str = None
):
    """
    IP addresses have number of relationships to other objects. This returns ALL objects that fit the relationship.

    The relationships are documented here:
    - comments: The comments for the IP address. Returns a list of comments.
    - communicating_files: Files that communicate with the IP address. Returns a list of files.
    - downloaded_files: Files downloaded from the IP address. VT Enterprise users only. Returns a list of files.
    - graphs: Graphs including the IP address. Returns a list of graphs.
    - historical_ssl_certificates: SSL certificates associated with the IP. Returns a list of SSL certificates.
    - historical_whois: WHOIS information for the IP address. Retrurns a list of Whois attributes.
    - related_comments: Community posted comments in the IP's related objects. Returns a list of comments.
    - related_references: Returns the references related to the IP address. Returns a list of References.
    - related_threat_actors: Threat actors related to the IP address. Returns a list of threat actors.
    - referrer_files: Files containing the IP address. Returns a list of Files.
    - resolutions: Resolves the IP addresses. Returns a list of resolutions.
    - urls: Returns a list of URLs related to the IP address. Returns a list of URLs.

    Args:
    - ip, string, required, IP address
    - relationship, string, required, Relationship name (see the list of items from above)
    - x-apikey, string, required, Your API key
    - limit, int32, optional, Maximum number of comments to retrieve
    - cursor, string, optional, Continuation cursor
    """


def vt_get_ip_address_report(ip: str, x_apikey: str):
    """
    Retrieve an IP address report. These reports condense all of the recent activity that VirusTotal has seen for the resource under consideration, as well as contextual information about it.
    This function specifically generates these reports using the IP address parameter.

    Args:
    - ip: string, required, IP address
    - x-apikey: string, required, Your API key
    """


def vt_add_votes_to_ip_address(ip: str, data: dict, x_apikey: str):
    """
    With this function you can post a vote for a given file. The body for the POST request must be the JSON representation of a vote object. Note however that you don't need to provide an ID for the object, as they are automatically generated for new votes. The verdict attribute must have be either harmless or malicious.

    Please ensure that the JSON object you provide conforms accurately to valid JSON standards.

    Args:
    - ip, string, required, IP address
    - data, json, Vote object
    - x-apikey, string, required, Your API key
    """


def vt_get_domain_report(domain: str, x_apikey: str):
    """
    Retrieves a domain report. These reports contain information regarding the domain itself that VirusTotal has collected.

    Args:
    - domain: string, required, Domain name
    - x-apikey: string, required, Your API key
    """


def vt_get_comments_on_ip_address(
    ip: str, x_apikey: str, limit: int = None, cursor: str = None
):
    """
    Retrieves the comments on a provided IP address. Returns a list of Comment objects.

    Args:
    - ip, string, required, IP address
    - x-apikey, string, required, Your API key
    - limit, int32, optional, Maximum number of comments to retrieve
    - cursor, string, optional, Continuation cursor
    """


def vt_add_comment_to_ip_address(ip: str, data: dict, x_apikey: str):
    """
    With this function you can post a comment for a given IP address. The body for the POST request must be the JSON representation of a comment object. Notice however that you don't need to provide an ID for the object, as they are automatically generated for new comments.
    However, please note that you will need to provide a valid data JSON for using this function.

    Any word starting with # in your comment's text will be considered a tag, and added to the comment's tag attribute.

    Returns a Comment object.

    Args:
    - ip: string, required, IP address
    - data: json, required, A comment object
    - x-apikey: string, required, Your API key
    """


def vt_get_object_descriptors_related_to_ip_address(
    ip: str, relationship: str, x_apikey: str, limit: int = None, cursor: str = None
):
    """
    This specifically returns related object's IDs (and context attributes, if any). Please note that this will not return all attributes.

    You are expected to provide the relationship to the object you're interested in. The valid relationships are as follows.

    The relationships are documented here:
    - comments: The comments for the IP address.
    - communicating_files: Files that communicate with the IP address.
    - downloaded_files: Files downloaded from the IP address. VT Enterprise users only.
    - graphs: Graphs including the IP address.
    - historical_ssl_certificates: SSL certificates associated with the IP.
    - historical_whois: WHOIS information for the IP address. Retrurns a list of Whois attributes.
    - related_comments: Community posted comments in the IP's related objects.
    - related_references: Returns the references related to the IP address.
    - related_threat_actors: Threat actors related to the IP address.
    - referrer_files: Files containing the IP address.
    - resolutions: Resolves the IP addresses.
    - urls: Returns a list of URLs related to the IP address.

    Here are some useful descriptions of the arguments in this API, with the format - name of this argument: type of the data, required or optional, description of this argument.
    - ip: string, required, IP address
    - relationship: string, required, Relationship name (see table)
    - x-apikey: string, required, Your API key
    - limit: int32, optional, Maximum number of comments to retrieve
    - cursor: string, optional, Continuation cursor
    """


def vt_get_objects_related_to_domain(
    domain: str, relationship: str, x_apikey: str, limit: int = None, cursor: str = None
):
    """
    Objects are a key concept in the VirusTotal API. Each object has an identifier and a type.
    Each object has an associated URL, and each domain is associated with objects.
    This function returns ALL of the objects related to the domain, based on the specified relationship.

    The following describe the valid relationship:
    - caa_records: Records CAA for the domain.
    - cname_records: Records CNAME for the domain.
    - comments: Community posted comments about the domain.
    - communicating_files: Files that communicate with the domain.
    - downloaded_files: Files downloaded from that domain.
    - graphs: All graphs that include the domain.
    - historical_ssl_certificates: SSL certificates associated with the domain.
    - historical_whois: WHOIS information for the domain.
    - immediate_parent: Domain's immediate parent.
    - mx_records: Records MX for the domain.
    - ns_records: Records NS for the domain.
    - parent: Domain's top parent.
    - referrer_files: Refers to any and all files that contain this domain.
    - related_comments: Community posted comments in the domain's related objects.
    - related_references: Refers to the References related to the domain.
    - related_threat_actors: Refers to the threat actors related to the domain. A list of Threat Actors.
    - resolutions: DNS resolutions for the domain.
    - soa_records: Records SOA for the domain.
    - siblings: Refers to the Domain's sibling domains.
    - subdomains: Refers to the Domain's subdomains.
    - urls: Refers to the URLs that contain this domain.
    - user_votes: Refers to the current user's votes.


    Args:
    - domain: string, required, Domain name
    - relationship, string, required, Relationship name (see table)
    - x-apikey, string, required, Your API key
    - limit, int32, optional, Maximum number of comments to retrieve
    - cursor, string, optional, Continuation cursor
    """


def vt_get_object_descriptors_related_to_domain(
    domain: str, relationship: str, x_apikey: str, limit: int = None, cursor: str = None
):
    """
    This specifically returns related object's IDs (and context attributes, if any). Please note that this will not return all attributes. This will return objects relating to a domain.

    - caa_records: Records CAA for the domain.
    - cname_records: Records CNAME for the domain.
    - comments: Community posted comments about the domain.
    - communicating_files: Files that communicate with the domain.
    - downloaded_files: Files downloaded from that domain.
    - graphs: All graphs that include the domain.
    - historical_ssl_certificates: SSL certificates associated with the domain.
    - historical_whois: WHOIS information for the domain.
    - immediate_parent: Domain's immediate parent.
    - mx_records: Records MX for the domain.
    - ns_records: Records NS for the domain.
    - parent: Domain's top parent.
    - referrer_files: Refers to any and all files that contain this domain.
    - related_comments: Community posted comments in the domain's related objects.
    - related_references: Refers to the References related to the domain.
    - related_threat_actors: Refers to the threat actors related to the domain. A list of Threat Actors.
    - resolutions: DNS resolutions for the domain.
    - soa_records: Records SOA for the domain.
    - siblings: Refers to the Domain's sibling domains.
    - subdomains: Refers to the Domain's subdomains.
    - urls: Refers to the URLs that contain this domain.
    - user_votes: Refers to the current user's votes.

    Args:
    - domain: string, required, Domain name
    - relationship: string, required, Relationship name (see table)
    - x-apikey: string, required, Your API key
    - limit: int32, optional, Maximum number of comments to retrieve
    - cursor: string, optional, Continuation cursor
    """


def vt_get_comments_on_domain(
    domain: str, x_apikey: str, limit: int = None, cursor: str = None
):
    """
    This function will retrieve comments on a specified domain.

    Args:
    - domain, string, required, Domain name
    - x-apikey, string, required, Your API key
    - limit, int32, optional, Maximum number of comments to retrieve
    - cursor, string, optional, Continuation cursor
    """


def vt_get_votes_on_ip_address(ip: str):
    """
    This function will retrieve votes on a provided IP address.

    Args:
    - ip: string, required, ip address
    """

raven/data/resources/virustotal_function_definitions.py
nexusflowai/NexusRaven
from typing import Callable, Dict, List

from dataclasses import dataclass

import json

# Necessary imports to successfully eval ToolLLM functions
from datetime import datetime  # pylint: disable=unused-import
from enum import Enum as ENUM  # pylint: disable=unused-import

from datasets import Dataset, load_dataset


@dataclass
class ToolLLMEvaluationDataHelper:
    hf_path: str
    standardized_queries_subset: str

    def __post_init__(self) -> None:
        raw_queries_subset = self.standardized_queries_subset.replace(
            "standardized", "raw"
        )
        self.dataset = load_dataset(
            path=self.hf_path,
            name=raw_queries_subset,
            split="train",
        )
        self.dataset = self.dataset.filter(
            lambda s: s == "toolllm", input_columns="dataset"
        )
        self.processed_dataset: Dataset | None = None

    def build_functions(self) -> Dict[str, Callable]:
        context_functions: List[List[str]] = []
        prompts: List[str] = []
        reference_function_calls: List[str] = []
        functions: Dict[str, Callable] = dict()
        for example in self.dataset:
            example = json.loads(example["query_dict"])
            function_strs = example["context"]

            function_str = None
            function_name = None
            function_names = []
            for function_str in function_strs:
                # Replace braces since they cause problems for langchain's formatting
                function_str = function_str.replace("{", "")
                function_str = function_str.replace("}", "")

                function_name = function_str[: function_str.find("(")].removeprefix(
                    "def "
                )

                # Remove the function body to use our own for evaluation
                function_str = function_str[: function_str.find("args_dict")]
                function_str = f"""
{function_str}return ("{function_name}", {{k: int(v) if isinstance(v, float) else v for k, v in locals().items()}})
"""

                exec(function_str)

                function_names.append(function_name)

            namespace = locals()
            new_functions = {n: namespace[n] for n in function_names}
            functions.update(new_functions)

            context_functions.append(function_names)

            prompts.append(example["Input"])
            reference_function_calls.append(example["Output"])

        self.processed_dataset = Dataset.from_dict(
            {
                "context_functions": context_functions,
                "prompt": prompts,
                "reference_function_call": reference_function_calls,
            }
        )

        return functions

    def get_eval_dataset(self) -> Dataset:
        if self.processed_dataset is None:
            self.build_functions()

        return self.processed_dataset

raven/data/toolllm_evaluation_data_utils.py
nexusflowai/NexusRaven
from typing import Any, Dict, List, Tuple

from dataclasses import dataclass

import os

from argparse import ArgumentParser

from collections import defaultdict

import re

import json

from datasets import Dataset, concatenate_datasets, load_dataset

from raven import ROOT_DIR


@dataclass
class ToolLLMFormatQueriesHelper:
    hf_path: str
    toolllm_queries_subset: str
    toolllm_api_list_subset: str
    standardized_queries_subset: str
    standardized_api_list_subset: str

    def run(self) -> None:
        api_list = load_dataset(
            self.hf_path, name=self.standardized_api_list_subset, split="train"
        )
        queries = load_dataset(
            self.hf_path, name=self.standardized_queries_subset, split="train"
        )

        api_list, queries = self.concatenate_toolllm(api_list, queries)

        python_functions = list(
            map(self.convert_to_python_function, api_list.to_iterable_dataset())
        )
        python_files = self.group_python_functions_by_name(python_functions)
        python_files = Dataset.from_list(python_files)

        json_function_descriptions = list(
            map(
                self.convert_to_json_function_description,
                api_list.to_iterable_dataset(),
            )
        )
        json_function_descriptions = self.group_json_functions_by_name(
            json_function_descriptions
        )
        json_function_descriptions = Dataset.from_list(json_function_descriptions)

        custom_files = Dataset.from_list(self.get_custom_files())

        json_query_dicts = list(
            map(self.convert_to_json_query, queries.to_iterable_dataset())
        )
        self.add_query_id(json_query_dicts)

        toolllm_queries = Dataset.from_list(json_query_dicts)
        toolllm_api_list = concatenate_datasets(
            [python_files, json_function_descriptions, custom_files]
        )

        toolllm_queries.push_to_hub(
            repo_id=self.hf_path,
            config_name=self.toolllm_queries_subset,
        )
        toolllm_api_list.push_to_hub(
            repo_id=self.hf_path,
            config_name=self.toolllm_api_list_subset,
        )

    def concatenate_toolllm(
        self, api_list: Dataset, queries: Dataset
    ) -> Tuple[Dataset, Dataset]:
        """
        ToolLLM data is not yet fully integrated into this repository, so here we take a chunk from the
        Nexusflow/NexusRaven_API_evaluation dataset
        """
        toolllm_api_list = load_dataset(
            path="Nexusflow/NexusRaven_API_evaluation",
            name="standardized_api_list",
            split="train",
        )
        toolllm_api_list = toolllm_api_list.filter(
            lambda s: s == "toolllm", input_columns="dataset"
        )
        api_list = api_list.filter(lambda s: s != "toolllm", input_columns="dataset")
        api_list = concatenate_datasets([api_list, toolllm_api_list])

        toolllm_queries = load_dataset(
            path="Nexusflow/NexusRaven_API_evaluation",
            name="standardized_queries",
            split="train",
        )
        toolllm_queries = toolllm_queries.filter(
            lambda s: s == "toolllm", input_columns="dataset"
        )
        queries = queries.filter(lambda s: s != "toolllm", input_columns="dataset")
        queries = concatenate_datasets([queries, toolllm_queries])

        return api_list, queries

    def convert_to_python_function(self, api_dict: Dict[str, Any]) -> Dict[str, str]:
        name = api_dict["name"]
        description = api_dict["description"]
        args_dicts = api_dict["args_dicts"]

        # Convert to lower since ToolLLM only accepts lowercase function names
        lowered_name = self.to_lower_str(name)
        description = description.replace(name, lowered_name)
        name = lowered_name
        for arg_dict in args_dicts:
            lowered_name = self.to_lower_str(arg_dict["name"])
            description = description.replace(arg_dict["name"], lowered_name)
            arg_dict["name"] = lowered_name

        args_dicts = sorted(args_dicts, key=lambda d: d["required"], reverse=True)

        args_strs = []
        for arg_dict in args_dicts:
            if arg_dict["required"]:
                default_str = ""
            else:
                default = arg_dict["default"]
                if isinstance(default, str) and default.startswith("JSON:"):
                    default = default.removeprefix("JSON:")
                    default = json.loads(default)
                default_str = f" = {default}"

            if arg_dict["type"] == "None":
                type_str = ""
            else:
                type_str = f": {arg_dict['type']}"

            args_str = f"{arg_dict['name']}{type_str}{default_str}"
            args_strs.append(args_str)

        function_str = f'def {name}({", ".join(args_strs)}):\n    """{description}"""\n    return "Is successful."'

        return {
            "dataset": api_dict["dataset"],
            "function_str": function_str,
        }

    def group_python_functions_by_name(
        self, function_dicts: List[Dict[str, str]]
    ) -> List[Dict[str, str]]:
        groups = defaultdict(list)
        for function_dict in function_dicts:
            groups[function_dict["dataset"]].append(function_dict["function_str"])

        python_files = []
        for group_name, group_function_strs in groups.items():
            group_function_str = "\n\n".join(group_function_strs)
            group_dict = {
                "file_path": f"data/toolenv/tools/Customized/{group_name}/api.py",
                "file_content": group_function_str,
            }
            python_files.append(group_dict)

        return python_files

    def convert_to_json_function_description(
        self, api_dict: Dict[str, Any]
    ) -> Dict[str, Any]:
        name = api_dict["name"]
        description = api_dict["description"]
        args_dicts = api_dict["args_dicts"]

        # Convert to lower since ToolLLM only accepts lowercase function names
        lowered_name = self.to_lower_str(name)
        description = description.replace(name, lowered_name)
        name = lowered_name
        for arg_dict in args_dicts:
            lowered_name = self.to_lower_str(arg_dict["name"])
            description = description.replace(arg_dict["name"], lowered_name)
            arg_dict["name"] = lowered_name

        required_parameters = []
        optional_parameters = []
        for arg_dict in args_dicts:
            tool_llm_arg_dict = arg_dict.copy()
            tool_llm_arg_dict.pop("required")

            required = arg_dict["required"]
            if required:
                list_to_append_to = required_parameters
            else:
                list_to_append_to = optional_parameters
            list_to_append_to.append(tool_llm_arg_dict)

        function_dict = {
            "name": name,
            "url": "",
            "description": description,
            "method": "",
            "required_parameters": required_parameters,
            "optional_parameters": optional_parameters,
        }

        return {
            "dataset": api_dict["dataset"],
            "function_dict": function_dict,
        }

    def group_json_functions_by_name(
        self, function_dicts: List[Dict[str, str]]
    ) -> List[Dict[str, str]]:
        groups = defaultdict(list)
        for function_dict in function_dicts:
            groups[function_dict["dataset"]].append(function_dict["function_dict"])

        json_function_dicts = []
        for group_name, group_function_dicts in groups.items():
            tool_dict = {
                "tool_description": group_name,
                "tool_name": group_name,
                "title": group_name,
                "standardized_name": group_name,
                "api_list": group_function_dicts,
            }

            json_function_dict = {
                "file_path": f"data/toolenv/tools/Customized/{group_name}.json",
                "file_content": json.dumps(tool_dict, indent=4),
            }
            json_function_dicts.append(json_function_dict)

        return json_function_dicts

    def get_custom_files(self) -> List[Dict[str, str]]:
        resources_path = os.path.join(ROOT_DIR, "data", "resources")
        rapidapi = {
            "file_path": "toolbench/inference/Downstream_tasks/rapidapi.py",
            "file_content": open(os.path.join(resources_path, "rapidapi.py")).read(),
        }
        server = {
            "file_path": "toolbench/inference/server.py",
            "file_content": open(os.path.join(resources_path, "server.py")).read(),
        }

        return [rapidapi, server]

    def convert_to_json_query(self, query_dict: Dict[str, Any]) -> Dict[str, str]:
        api_list = []
        for context_function in query_dict["context_functions"]:
            context_function_dict = {
                "api_name": self.to_lower_str(context_function),
                "category_name": "Customized",
                "tool_name": query_dict["dataset"],
            }
            api_list.append(context_function_dict)

        return {
            "api_list": api_list,
            "query": query_dict["prompt"],
        }

    def add_query_id(self, json_query_dicts: List[Dict[str, str]]) -> None:
        query_id = 80000
        for d in json_query_dicts:
            d["query_id"] = query_id
            query_id += 1

    @staticmethod
    def to_lower_str(s: str) -> str:
        # `lowerCamelCase` becomes `lowerxxxcamelxxxcase`
        s = re.sub("([A-Z])", r"xxx\1", s)
        s = s.lower()
        return s


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--hf_path", type=str, required=True)
    parser.add_argument("--toolllm_api_list_subset", type=str, required=True)
    parser.add_argument("--toolllm_queries_subset", type=str, required=True)
    parser.add_argument("--standardized_api_list_subset", type=str, required=True)
    parser.add_argument("--standardized_queries_subset", type=str, required=True)
    args = parser.parse_args()

    ToolLLMFormatQueriesHelper(**vars(args)).run()

raven/data/upload_queries_in_toolllm_format.py
nexusflowai/NexusRaven
from dataclasses import dataclass

from argparse import ArgumentParser

import os

import json

from datasets import Dataset, concatenate_datasets

from raven import ROOT_DIR


@dataclass
class RawQueriesHelper:
    hf_path: str
    subset: str

    def run(self) -> None:
        d = concatenate_datasets(
            [
                self.upload_local_queries("cve_cpe"),
                self.upload_local_queries("emailrep"),
                self.upload_local_queries("virustotal"),
                self.upload_local_queries("toolalpaca"),
                self.upload_local_queries("toolllm"),
            ]
        )
        d.push_to_hub(
            repo_id=self.hf_path,
            config_name=self.subset,
        )

    def upload_local_queries(self, dataset: str) -> Dataset:
        with open(
            os.path.join(ROOT_DIR, "data", "resources", f"{dataset}_queries.json")
        ) as f:
            content = json.load(f)

        dicts = list(map(json.dumps, content))
        return Dataset.from_dict(
            {
                "dataset": [dataset] * len(dicts),
                "query_dict": dicts,
            }
        )


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--hf_path", type=str, required=True)
    parser.add_argument("--subset", type=str, required=True)
    args = parser.parse_args()

    RawQueriesHelper(**vars(args)).run()

raven/data/upload_raw_queries.py
nexusflowai/NexusRaven
from typing import Any, Dict, List

from dataclasses import dataclass

from argparse import ArgumentParser

import os

import ast

import re

from datasets import Dataset, concatenate_datasets, load_dataset

from raven.utils import build_functions


@dataclass
class UploadStandardizedAPIListHelper:
    hf_path: str
    subset: str

    def run(self) -> None:
        d = concatenate_datasets(
            [
                self.cve_cpe(),
                self.emailrep(),
                self.virustotal(),
                self.toolalpaca(),
            ]
        )

        d: Dataset = d.map(
            function=self.standardize_arguments,
            input_columns="args_dicts",
            desc="Standardizing arguments",
        )

        build_functions(d)

        d.push_to_hub(
            repo_id=self.hf_path,
            config_name=self.subset,
        )

    def standardize_arguments(
        self, ds: List[Dict[str, Any]]
    ) -> Dict[str, List[Dict[str, Any]]]:
        for d in ds:
            d["name"] = d["name"].replace("-", "_")  # Argument names must be pythonic

            if d["type"] in ["str", "list", "bool", "int", "dict", "None"]:
                pass
            elif "list" in d["type"].lower():
                d["type"] = "list"
            elif "array" in d["type"].lower():
                d["type"] = "list"
            elif "bool" in d["type"].lower():
                d["type"] = "bool"
            elif "int" in d["type"].lower():
                d["type"] = "int"
            elif "str" in d["type"].lower():
                d["type"] = "str"
            elif "json" in d["type"].lower():
                d["type"] = "str"
            elif "enum" in d["type"].lower():
                d["type"] = "str"
            elif "datetime" in d["type"].lower():
                d["type"] = "str"
            else:
                raise ValueError(f"Unrecognized argument type `{d['type']}`")

        return {"args_dicts": ds}

    def cve_cpe(self) -> Dataset:
        p = os.path.split(__file__)[0]
        function_definitions_path = os.path.join(
            p, "resources", "cve_cpe_function_definitions.py"
        )
        with open(function_definitions_path) as f:
            content = f.read()

        content = ast.parse(content).body

        function_dicts: List[Dict[str, Any]] = []
        function: ast.FunctionDef
        for function in content:
            description = function.body[0].value.value

            args_str = description.split("Args:")[1].strip()
            args_strs = filter(None, args_str.split("\n"))
            args_dicts = []
            for arg_str in args_strs:
                name_type_str = arg_str[: arg_str.find(":")].strip()
                name = name_type_str.split()[0].strip()
                description = arg_str[arg_str.find(":") :].strip()
                arg_dict = {
                    "name": name,
                    "type": "None",
                    "required": False,  # All optional
                    "description": description,
                    "default": "None",
                }
                args_dicts.append(arg_dict)

            function_dict = {
                "dataset": "cve_cpe",
                "name": function.name,
                "description": function.body[0].value.value,
                "args_dicts": args_dicts,
            }
            function_dicts.append(function_dict)

        return Dataset.from_list(function_dicts)

    def emailrep(self) -> Dataset:
        """
        Function structure from https://github.com/sublime-security/emailrep.io-python/blob/9e187c9f8d1ba25fae0e32fc92e4ccc380b77933/emailrep/__init__.py#L6

        Argument descriptions
        - emailrep_post https://docs.sublimesecurity.com/reference/get_-email
        - emailrep_report https://docs.sublimesecurity.com/reference/post_report

        We additionally add a required `key` argument to the `emailrep_report` arguments since it's required to create a report.
        We remove the `description` and `tags` arguments since they are arbitrary.
        We remove the `timestamp` and `expires` args since they are never used.

        """

        emailrep_post = {
            "dataset": "emailrep",
            "name": "emailrep_post",
            "description": """
Query an email address

Parameters
email (string, required): Email address being queried
""",
            "args_dicts": [
                {
                    "name": "email",
                    "description": "Email address being queried",
                    "type": "string",
                    "default": "None",
                    "required": True,
                },
            ],
        }

        emailrep_report = {
            "dataset": "emailrep",
            "name": "emailrep_report",
            "description": """
Reports an email address. Date of malicious activity defaults to the current time unless otherwise specified.

Parameters:
email (string, required): Email address being reported.
key (string, required): The API key of the user.
""",
            "args_dicts": [
                {
                    "name": "email",
                    "description": "Email address being reported.",
                    "type": "string",
                    "default": "None",
                    "required": True,
                },
                {
                    "name": "key",
                    "description": "The API key of the user.",
                    "type": "string",
                    "default": "None",
                    "required": True,
                },
            ],
        }

        return Dataset.from_list([emailrep_post, emailrep_report])

    def virustotal(self) -> Dataset:
        """
        All function documentation is sourced from https://developers.virustotal.com/reference/overview
        since the existing VirusTotal Python client does not provide function definitions https://github.com/VirusTotal/vt-py/. vt-py just calls into the overall VirusTotal http endpoint

        """
        p = os.path.split(__file__)[0]
        function_definitions_path = os.path.join(
            p, "resources", "virustotal_function_definitions.py"
        )
        with open(function_definitions_path) as f:
            content = f.read()

        content = ast.parse(content).body

        pattern = r"-\s(\w.+):\s(\w+),\s(\w+),\s(.+)"
        pattern = re.compile(pattern)

        function_dicts: List[Dict[str, Any]] = []
        function: ast.FunctionDef
        for function in content:
            args_list = list(function.args.args)
            defaults = list(function.args.defaults)
            defaults = [None] * (len(args_list) - len(defaults)) + defaults

            args_dicts = []
            for arg, default in zip(args_list, defaults):
                arg_dict = {
                    "name": arg.arg,
                    "type": arg.annotation.id,
                    "required": default is None,
                    "description": "",
                    "default": "None" if default is None else default.value,
                }
                args_dicts.append(arg_dict)

            function_dict = {
                "dataset": "virustotal",
                "name": function.name,
                "description": function.body[0].value.value,
                "args_dicts": args_dicts,
            }
            function_dicts.append(function_dict)

        return Dataset.from_list(function_dicts)

    def toolalpaca(self) -> Dataset:
        d = load_dataset("Nexusflow/toolalpaca_simulated", split="train")
        all_docstrings: List[str] = [_d["docstring"] for _d in d["turns_1"]]

        func_pattern = re.compile("<func_start>(.*?)<func_end>", re.S)
        docstring_pattern = re.compile("<docstring_start>(.*?)<docstring_end>", re.S)

        function_strs = []
        for docstrings_to_parse in all_docstrings:
            func_signatures = re.findall(func_pattern, docstrings_to_parse)
            docstrings = re.findall(docstring_pattern, docstrings_to_parse)
            for func_signature, docstring in zip(func_signatures, docstrings):
                function_str = f"{func_signature}:\n    {docstring.strip()}"
                function_strs.append(function_str)

        function_strs = list(dict.fromkeys(function_strs))  # dedupe

        function_dicts = []
        for function_str in function_strs:
            function: ast.FunctionDef = ast.parse(function_str).body[0]
            description = function.body[0].value.value
            description = description.replace("{", "").replace("}", "")

            args_dicts = []
            for arg in function.args.args:
                if isinstance(arg.annotation, ast.Name):
                    type_str = arg.annotation.id
                else:
                    type_str = arg.annotation.value.id
                arg_dict = {
                    "name": arg.arg,
                    "description": "",
                    "type": type_str,
                    "default": "None",  # No defaults
                    "required": True,
                }
                args_dicts.append(arg_dict)

            function_dict = {
                "dataset": "toolalpaca",
                "name": function.name,
                "description": description.strip(),
                "args_dicts": args_dicts,
            }
            function_dicts.append(function_dict)
        return Dataset.from_list(function_dicts)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--hf_path", type=str, required=True)
    parser.add_argument("--subset", type=str, required=True)
    args = parser.parse_args()

    UploadStandardizedAPIListHelper(**vars(args)).run()

raven/data/upload_standardized_api_list.py
nexusflowai/NexusRaven
from typing import Any, Dict, List

from dataclasses import dataclass

from collections import defaultdict

from argparse import ArgumentParser

import re

import json

from datasets import load_dataset

from raven.utils import build_functions, parse_function_call_to_name_and_args


@dataclass
class StandardizedQueriesHelper:
    hf_path: str
    standardized_queries_subset: str
    raw_queries_subset: str
    standardized_api_list_subset: str

    def run(self) -> None:
        api_list = load_dataset(
            self.hf_path, name=self.standardized_api_list_subset, split="train"
        )
        dataset_to_context_functions = defaultdict(list)
        for e in api_list:
            dataset_to_context_functions[e["dataset"]].append(e["name"])

        raw_queries = load_dataset(
            self.hf_path, name=self.raw_queries_subset, split="train"
        )
        # Filter out ToolLLM examples since those are a special case handled later on
        raw_queries = raw_queries.filter(
            lambda s: s != "toolllm", input_columns="dataset"
        )
        standardized_queries = raw_queries.map(
            function=self._map_single,
            input_columns=["dataset", "query_dict"],
            remove_columns=raw_queries.column_names,
            fn_kwargs={"dataset_to_context_functions": dataset_to_context_functions},
        )

        functions = build_functions(api_list)
        for standardized_query in standardized_queries:
            reference_function_name = standardized_query["python_function_name"]
            reference_input_args_dict = json.loads(
                standardized_query["python_args_dict"]
            )

            try:
                functions[reference_function_name](**reference_input_args_dict)
            except Exception as e:
                print(reference_function_name, reference_input_args_dict)
                raise e

        standardized_queries.push_to_hub(
            repo_id=self.hf_path,
            config_name=self.standardized_queries_subset,
        )

    def _map_single(
        self,
        dataset: str,
        query_dict: str,
        dataset_to_context_functions: Dict[str, List[str]],
    ) -> Dict[str, str]:
        query_dict: Dict[str, Any] = json.loads(query_dict)
        match dataset:
            case "cve_cpe":
                d = self.cve_cpe(query_dict)
            case "emailrep" | "virustotal":
                context_functions = dataset_to_context_functions[dataset]
                d = self.parse_input_output(query_dict, context_functions)
            case "toolalpaca":
                d = self.toolalpaca(query_dict)
            case _:
                raise ValueError(f"Unrecognized dataset `{dataset}`")

        d["dataset"] = dataset
        return d

    def cve_cpe(self, query_dict: Dict[str, Any]) -> Dict[str, str]:
        function = query_dict["reference"]
        function = function.removeprefix("r = ")
        function = function.replace("nvdlib.", "")
        query_dict["reference"] = function

        context_functions = ["searchCVE", "searchCPE"]
        return self.parse_input_output(query_dict, context_functions)

    def parse_input_output(
        self, query_dict: Dict[str, Any], context_functions: List[str]
    ) -> Dict[str, str]:
        function_name, _, args_dict = parse_function_call_to_name_and_args(
            query_dict["reference"]
        )
        d = {
            "prompt": query_dict["query"],
            "python_function_name": function_name,
            "python_args_dict": json.dumps(args_dict),
            "context_functions": context_functions,
        }
        return d

    def toolalpaca(self, query_dict: Dict[str, Any]) -> Dict[str, str]:
        query_dict = query_dict["turns_1"]  # Just the first turn

        # These function names must match the API name later on
        reference = json.loads(query_dict["golden_answer"])[0]
        function_name, _, args_dict = parse_function_call_to_name_and_args(reference)

        func_pattern = re.compile("<func_start>def (.*?)<func_end>", re.S)
        matches = func_pattern.findall(query_dict["docstring"])
        context_functions = []
        for match in matches:
            paren_idx = match.find("(")
            match = match[:paren_idx]
            context_functions.append(match)

        return {
            "prompt": query_dict["instruction"],
            "python_function_name": function_name,
            "python_args_dict": json.dumps(args_dict),
            "context_functions": context_functions,
        }


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--hf_path", type=str, required=True)
    parser.add_argument("--standardized_queries_subset", type=str, required=True)
    parser.add_argument("--raw_queries_subset", type=str, required=True)
    parser.add_argument("--standardized_api_list_subset", type=str, required=True)
    args = parser.parse_args()

    StandardizedQueriesHelper(**vars(args)).run()

raven/data/upload_standardized_queries.py
nexusflowai/NexusRaven
import langchain


langchain.debug = False

raven/eval/__init__.py
nexusflowai/NexusRaven
from typing import Any, Callable, Dict, List

from dataclasses import dataclass

import json

from os import environ

from string import ascii_lowercase

import argparse

from unittest.mock import MagicMock

from datasets import Dataset, load_dataset

from langchain.tools.base import StructuredTool
from langchain.llms import OpenAI, HuggingFaceTextGenInference, OpenAIChat
from langchain.chat_models import ChatOpenAI
from langchain.agents import (
    AgentExecutor,
    AgentType,
    LLMSingleActionAgent,
    initialize_agent,
)
from langchain.chains import LLMChain

from raven.utils import build_functions, parse_function_call_to_name_and_args
from raven.eval.raven_utils import (
    RavenOutputParser,
    RavenPromptTemplate,
    RAVEN_PROMPT,
)
from raven.data.toolllm_evaluation_data_utils import ToolLLMEvaluationDataHelper


@dataclass
class Evaluator:
    task_name: str
    llm_name: str
    agent_name: str
    hf_path: str
    standardized_queries_subset: str
    standardized_api_list_subset: str
    inference_server_url: str | None = None

    def __post_init__(self) -> None:
        self.toolllm_helper = ToolLLMEvaluationDataHelper(
            hf_path=self.hf_path,
            standardized_queries_subset=self.standardized_queries_subset,
        )
        self.agent: Any | None = None

    def run(self) -> None:
        functions = self.build_functions()
        locals().update(functions)

        llm = self.build_llm()
        tools = {
            function_name: StructuredTool.from_function(function, return_direct=True)
            for function_name, function in functions.items()
        }
        agent = self.build_agent(llm, list(tools.values()))

        eval_dataset = self.get_eval_dataset()
        accuracies = []

        for i, sample in enumerate(eval_dataset):
            context_tools = [tools[k] for k in sample["context_functions"]]
            agent = self.build_agent(llm, context_tools)

            prompt = sample["prompt"]
            error_message = None
            output = None
            original_output = None
            try:
                original_output = agent.run(prompt)
                if isinstance(original_output, tuple):  # Successful call
                    output = original_output
                elif isinstance(
                    original_output, str
                ):  # Model returned execution string
                    (
                        function_name,
                        args,  # pylint: disable=redefined-outer-name
                        kwargs,
                    ) = parse_function_call_to_name_and_args(original_output)
                    output = functions[function_name](*args, **kwargs)
                else:
                    raise ValueError(
                        f"Unknown output type `{original_output}` of type {type(original_output)}"
                    )
            except Exception as e:  # pylint: disable=broad-exception-caught
                error_message = f"Original output\n\t{original_output}\nOutput\n\t{output}\nError message\n\t{str(e)}\n\n"
                output = (None, None)

            predicted_function_name, predicted_args_dict = output

            if "reference_function_call" in sample:
                reference_function_call = sample["reference_function_call"]
                reference_function_name, reference_args_dict = eval(
                    reference_function_call
                )
            else:
                reference_function_name = sample["python_function_name"]
                reference_input_args_dict = json.loads(sample["python_args_dict"])
                _, reference_args_dict = functions[reference_function_name](
                    **reference_input_args_dict
                )

            function_name_match = predicted_function_name == reference_function_name
            args_dict_match = predicted_args_dict == reference_args_dict
            accuracy = function_name_match and args_dict_match

            accuracy_str = "CORRECT" if accuracy else "WRONG"
            error_message_str = "" if error_message is None else error_message
            print(
                f"Example {i+1} / {len(eval_dataset)}\n\nPrompt\n\t{prompt}\n\nReference\n\t{reference_function_name}, {reference_args_dict}\n\nPrediction\n\t{predicted_function_name}, {predicted_args_dict}\n\n{accuracy_str}\n\n{error_message_str}{'-' * 80}\n"
            )

            accuracies.append(accuracy)

        accuracy = 100 * sum(accuracies) / len(eval_dataset)
        print(f"Accuracy: {accuracy:.3f}%")

    def get_eval_dataset(self) -> Dataset:
        if self.task_name == "toolllm":
            return self.toolllm_helper.get_eval_dataset()

        d = load_dataset(
            path=self.hf_path,
            name=self.standardized_queries_subset,
            split="train",
        )
        d = d.filter(
            lambda dataset: dataset == self.task_name,
            input_columns="dataset",
        )
        assert (
            len(d) > 0
        ), f"Unknown task `{self.task_name}`. Available tasks: {list(dict.fromkeys(d['dataset']))}"

        return d

    def build_llm(self) -> Any:
        if "openai" in self.llm_name:
            api_key = environ.get("OPENAI_API_KEY", None)
            if api_key is None:
                api_key = input(
                    f"Please input your OpenAI API key to use the `{self.llm_name}` model: "
                )
                environ["OPENAI_API_KEY"] = api_key

        gpt_params = {
            "temperature": 0,
            "verbose": True,
        }
        hf_params = {
            "max_new_tokens": 400,
            "do_sample": False,
            "inference_server_url": self.inference_server_url,
            "temperature": 0.001,
        }
        match self.llm_name:
            case "openai_gpt4":
                return OpenAIChat(**gpt_params, model_name="gpt-4")
            case "chat_openai_gpt4":
                return ChatOpenAI(**gpt_params, model_name="gpt-4")
            case "openai_gpt3.5":
                return OpenAIChat(**gpt_params, model_name="gpt-3.5-turbo-16k")
            case "chat_openai_gpt3.5":
                return ChatOpenAI(**gpt_params, model_name="gpt-3.5-turbo-16k")
            case "openai_gpt3.5_instruct":
                return OpenAI(**gpt_params, model_name="gpt-3.5-turbo-instruct")
            case "nexusraven" | "prototype":
                return HuggingFaceTextGenInference(**hf_params)
            case _:
                return self.llm_name

    def build_functions(self) -> Dict[str, Callable]:
        if self.task_name == "toolllm":
            return self.toolllm_helper.build_functions()

        d = load_dataset(
            path=self.hf_path,
            name=self.standardized_api_list_subset,
            split="train",
        )
        d = d.filter(
            lambda dataset: dataset == self.task_name,
            input_columns="dataset",
        )
        return build_functions(d)

    def build_agent(self, llm: Any, tools: List[StructuredTool]) -> Any:
        match self.agent_name:
            case "OPENAI_FUNCTIONS":
                return initialize_agent(
                    tools,
                    llm,
                    agent=AgentType.OPENAI_FUNCTIONS,
                    verbose=True,
                    max_iterations=1,
                )
            case "SIMPLE":
                return initialize_agent(
                    tools,
                    llm,
                    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
                    verbose=True,
                    max_iterations=1,
                )
            case "SIMPLE_NONCHAT":
                return initialize_agent(
                    tools,
                    llm,
                    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                    verbose=True,
                    max_iterations=1,
                )
            case "NEXUSRAVEN":
                raven_prompt = RavenPromptTemplate(
                    template=RAVEN_PROMPT, tools=tools, input_variables=["input"]
                )
                llm_chain = LLMChain(llm=llm, prompt=raven_prompt)
                output_parser = RavenOutputParser()
                agent = LLMSingleActionAgent(
                    llm_chain=llm_chain,
                    output_parser=output_parser,
                    stop=["\nReflection:"],
                    allowed_tools=tools,
                )
                agent_chain = AgentExecutor.from_agent_and_tools(
                    agent=agent, tools=tools, verbose=True
                )
                return agent_chain
            case "TOOLLLM":
                return self._build_toolllm_agent(llm, tools)
            case "TOOLALPACA":
                return self._build_toolalpaca_agent(llm, tools)
            case _:
                raise KeyError(f"Invalid agent_name `{self.agent_name}`")

    def _build_toolllm_agent(self, llm: Any, tools: List[StructuredTool]) -> Any:
        if self.agent is not None:
            return self.agent

        responses = load_dataset(
            self.hf_path,
            name=llm,
            split="train",
        )

        prompt_to_response = dict()
        for response in responses["response"]:
            prompt = response[0]["query"]
            response = response[1]["function_call"]
            prompt_to_response[prompt] = response

        agent = MagicMock()

        def run(prompt: str) -> str:
            return prompt_to_response.get(prompt, None)

        agent.run = run

        self.agent = agent
        return agent

    def _build_toolalpaca_agent(self, llm: Any, tools: List[StructuredTool]) -> Any:
        """
        TODO this is currently hardcoded to a specific dataset path
        """

        if self.agent is not None:
            return self.agent

        responses = load_dataset(
            "Nexusflow/toolalpaca_eval",
            split="train",
        )
        responses = responses.filter(
            lambda s: s == llm, input_columns="toolalpaca_model"
        )

        def to_upper(s: str) -> str:
            for c in ascii_lowercase:
                s = s.replace(f"_{c}", c.upper())
            return s

        if self.task_name in ["virustotal", "emailrep"]:
            to_upper = lambda s: s  # pylint: disable=unnecessary-lambda-assignment

        prompt_to_response = dict()
        for d in responses:
            prompt = d["prompt"]

            output = d["toolalpaca_response"]
            if len(output) == 0:
                continue

            function_name, args_dict = output[0]
            function_name = function_name.removeprefix("Action: ")
            function_name_mapping = {
                "search_cve": "searchCVE",
                "search_cpe": "searchCPE",
            }
            function_name = function_name_mapping.get(function_name, function_name)

            args_dict = args_dict.removeprefix("Action Input: ")
            try:
                args_dict = eval(args_dict)
            except:  # pylint: disable=bare-except
                # Model output is not a valid python object
                continue

            args_strs = []
            for arg_name, arg_value in args_dict.items():
                if isinstance(arg_value, str):
                    arg_value_str = f'"{arg_value}"'
                else:
                    arg_value_str = arg_value

                arg_name = to_upper(arg_name)
                arg_str = f"{arg_name}={arg_value_str}"
                args_strs.append(arg_str)

            args_str = ", ".join(args_strs)

            function_str = f"{function_name}({args_str})"

            prompt_to_response[prompt] = function_str

        agent = MagicMock()

        def run(prompt: str) -> str:
            return prompt_to_response.get(prompt, None)

        agent.run = run

        self.agent = agent

        return agent


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--task_name", type=str, required=True)
    parser.add_argument("--llm_name", type=str, required=True)
    parser.add_argument("--agent_name", type=str, required=True)
    parser.add_argument("--hf_path", type=str, required=True)
    parser.add_argument("--standardized_queries_subset", type=str, required=True)
    parser.add_argument("--standardized_api_list_subset", type=str, required=True)
    parser.add_argument("--inference_server_url", type=str)
    args = parser.parse_args()

    Evaluator(**vars(args)).run()

raven/eval/evaluator.py
nexusflowai/NexusRaven
from typing import List, Union

from langchain.agents import Tool, AgentOutputParser
from langchain.schema import AgentAction, AgentFinish, OutputParserException
from langchain.prompts import StringPromptTemplate


RAVEN_PROMPT = """
{raven_tools}
User Query: Question: {input}

Please pick a function from the above options that best answers the user query and fill in the appropriate arguments.<human_end>"""


def construct_raven_prompt(raven_tools):
    return RAVEN_PROMPT.format(raven_tools=raven_tools)


# Set up a prompt template
class RavenPromptTemplate(StringPromptTemplate):
    # The template to use
    template: str
    # The list of tools available
    tools: List[Tool]

    def format(self, **kwargs) -> str:
        prompt = "<human>:\n"
        for tool in self.tools:
            func_signature, func_docstring = tool.description.split(" - ", 1)
            prompt += f'\nOPTION:\n<func_start>def {func_signature}<func_end>\n<docstring_start>\n"""\n{func_docstring}\n"""\n<docstring_end>\n'
        kwargs["raven_tools"] = prompt
        return self.template.format(**kwargs).replace("{{", "{").replace("}}", "}")


class RavenOutputParser(AgentOutputParser):
    def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
        # Check if agent should finish
        if not "Initial Answer:" in text:
            raise OutputParserException(f"Could not parse LLM output: `{text}`")

        output = text.strip().split("\n")[1].replace("Initial Answer: ", "").strip()
        return AgentFinish(return_values={"output": output}, log=text)

raven/eval/raven_utils.py
nexusflowai/NexusRaven
from typing import Dict, List

from dataclasses import dataclass

from argparse import ArgumentParser

import os

import re

import json

from datasets import Dataset, load_dataset

from raven import CACHE_DIR


@dataclass
class RunToolLLM:
    hf_path: str
    toolllm_api_list_subset: str
    toolllm_queries_subset: str
    toolllm_outputs_subset: str
    tool_name: str | None = None

    def run(self) -> None:
        toolllm_path = os.path.join(CACHE_DIR, "ToolBench")
        setup_command = f"""
rm -rf {toolllm_path}/output* {toolllm_path}/results.txt
git clone https://github.com/OpenBMB/ToolBench {toolllm_path}
cd {toolllm_path}
pip -q -q -q install -r requirements.txt --user
wget -nc 'https://docs.google.com/uc?export=download&id=1Vis-RxBstXLKC1W1agIQUJNuumPJrrw0&confirm=yes' -O data.zip
unzip -n data.zip
"""
        os.system(setup_command)

        files_to_move = load_dataset(
            path=self.hf_path,
            name=self.toolllm_api_list_subset,
            split="train",
        )
        for file_dict in files_to_move:
            file_path = os.path.join(toolllm_path, file_dict["file_path"])
            os.makedirs(os.path.dirname(file_path), exist_ok=True)
            with open(file_path, "w") as f:
                f.write(file_dict["file_content"])

        query_dataset = load_dataset(
            path=self.hf_path,
            name=self.toolllm_queries_subset,
            split="train",
        )
        if self.tool_name is not None:
            query_dataset = query_dataset.filter(
                lambda l: l[0]["tool_name"] == self.tool_name,
                input_columns="api_list",
            )
        query_path = os.path.join(toolllm_path, "query.json")
        with open(query_path, "w") as f:
            f.write(json.dumps(list(query_dataset.to_iterable_dataset()), indent=4))

        os.system(
            f"""
cd {toolllm_path}
export PYTHONPATH=./
python toolbench/inference/qa_pipeline.py \
    --tool_root_dir data/toolenv/tools/ \
    --backbone_model toolllama \
    --model_path ToolBench/ToolLLaMA-7b \
    --max_observation_length 1024 \
    --observ_compress_method truncate \
    --method DFS_woFilter_w2 \
    --input_query_file query.json \
    --output_answer_file ./output_query.json \
    --api_customization
"""
        )

        results_path = os.path.join(toolllm_path, "results.txt")
        responses = self.parse_results_file_into_responses(results_path)
        responses = Dataset.from_dict({"response": responses})
        responses.push_to_hub(
            repo_id=self.hf_path,
            config_name=self.toolllm_outputs_subset,
        )

    def parse_results_file_into_responses(self, results_path: str) -> List[List[Dict]]:
        with open(results_path) as f:
            lines = map(str.strip, f.readlines())

        curr: List[Dict] = []
        res: List[List[Dict]] = []
        for line in lines:
            if line == "========":
                if len(curr) > 1:
                    res.append(curr)
                curr = []
            else:
                d = json.loads(line)
                if "function_call" in d:
                    d["function_call"] = self.to_upper_str(d["function_call"])

                curr.append(d)

        res.append(curr)

        return res

    def to_upper_str(self, s: str) -> str:
        # `lowerxxxcamelxxxcase` becomes `lowerCamelCase`
        s = re.sub("xxx([a-z])", lambda x: x.group(1).upper(), s)
        return s


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--hf_path", type=str, required=True)
    parser.add_argument("--toolllm_api_list_subset", type=str, required=True)
    parser.add_argument("--toolllm_queries_subset", type=str, required=True)
    parser.add_argument("--toolllm_outputs_subset", type=str, required=True)
    parser.add_argument("--tool_name", type=str)
    args = parser.parse_args()

    RunToolLLM(**vars(args)).run()

raven/eval/run_toolllm.py
nexusflowai/NexusRaven
from typing import Any, Callable, Dict, List, Tuple

import ast

import json

from datasets import Dataset


def build_functions(d: Dataset) -> Dict[str, Callable]:
    function_names = []
    for function_dict in d:
        name = function_dict["name"]
        description = function_dict["description"]
        args_dicts = function_dict["args_dicts"]

        args_dicts = sorted(args_dicts, key=lambda d: d["required"], reverse=True)

        args_strs = []
        for arg_dict in args_dicts:
            if arg_dict["required"]:
                default_str = ""
            else:
                default = arg_dict["default"]
                if isinstance(default, str) and default.startswith("JSON:"):
                    default = default.removeprefix("JSON:")
                    default = json.loads(default)
                default_str = f" = {default}"

            if arg_dict["type"] == "None":
                type_str = ""
            else:
                type_str = f": {arg_dict['type']}"

            args_str = f"{arg_dict['name']}{type_str}{default_str}"
            args_strs.append(args_str)

        function_str = f'def {name}({", ".join(args_strs)}):\n    """{description}"""\n    return ("{name}", locals())'
        exec(function_str)
        function_names.append(name)

    namespace = locals()
    functions = {f: namespace[f] for f in function_names}
    return functions


def parse_function_call_to_name_and_args(
    function: str,
) -> Tuple[str, List[str], Dict[str, Any]]:
    function = ast.parse(function).body[0]
    function_name = function.value.func.id
    keywords = function.value.keywords
    keywords = {k.arg: ast.literal_eval(k.value) for k in keywords}

    args = [ast.literal_eval(arg) for arg in function.value.args]

    """
    We use integers only for evaluation data since floats are tricky to evaluate
    e.g. the string representation of "3.33" may be "3.329999999" which yields an exact match of 0
    """
    keywords = {k: int(v) if isinstance(v, float) else v for k, v in keywords.items()}

    return function_name, args, keywords

raven/utils.py
nexusflowai/NexusRaven
from typing import List, Literal, Union

import math

from langchain.tools.base import StructuredTool
from langchain.agents import (
    Tool,
    AgentExecutor,
    LLMSingleActionAgent,
    AgentOutputParser,
)
from langchain.schema import AgentAction, AgentFinish, OutputParserException
from langchain.prompts import StringPromptTemplate
from langchain.llms import HuggingFaceTextGenInference
from langchain.chains import LLMChain


##########################################################
# Step 1: Define the functions you want to articulate. ###
##########################################################


def calculator(
    input_a: float,
    input_b: float,
    operation: Literal["add", "subtract", "multiply", "divide"],
):
    """
    Computes a calculation.

    Args:
    input_a (float) : Required. The first input.
    input_b (float) : Required. The second input.
    operation (string): The operation. Choices include: add to add two numbers, subtract to subtract two numbers, multiply to multiply two numbers, and divide to divide them.
    """
    match operation:
        case "add":
            return input_a + input_b
        case "subtract":
            return input_a - input_b
        case "multiply":
            return input_a * input_b
        case "divide":
            return input_a / input_b


def cylinder_volume(radius, height):
    """
    Calculate the volume of a cylinder.

    Parameters:
    - radius (float): The radius of the base of the cylinder.
    - height (float): The height of the cylinder.

    Returns:
    - float: The volume of the cylinder.
    """
    if radius < 0 or height < 0:
        raise ValueError("Radius and height must be non-negative.")

    volume = math.pi * (radius**2) * height
    return volume


#############################################################
# Step 2: Let's define some utils for building the prompt ###
#############################################################


RAVEN_PROMPT = """
{raven_tools}
User Query: Question: {input}

Please pick a function from the above options that best answers the user query and fill in the appropriate arguments.<human_end>"""


# Set up a prompt template
class RavenPromptTemplate(StringPromptTemplate):
    # The template to use
    template: str
    # The list of tools available
    tools: List[Tool]

    def format(self, **kwargs) -> str:
        prompt = "<human>:\n"
        for tool in self.tools:
            func_signature, func_docstring = tool.description.split(" - ", 1)
            prompt += f'\nOPTION:\n<func_start>def {func_signature}<func_end>\n<docstring_start>\n"""\n{func_docstring}\n"""\n<docstring_end>\n'
        kwargs["raven_tools"] = prompt
        return self.template.format(**kwargs).replace("{{", "{").replace("}}", "}")


class RavenOutputParser(AgentOutputParser):
    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
        # Check if agent should finish
        if "Initial Answer:" in llm_output:
            return AgentFinish(
                return_values={
                    "output": llm_output.strip()
                    .split("\n")[1]
                    .replace("Initial Answer: ", "")
                    .strip()
                },
                log=llm_output,
            )
        else:
            raise OutputParserException(f"Could not parse LLM output: `{llm_output}`")


##################################################
# Step 3: Build the agent with these utilities ###
##################################################


inference_server_url = "<YOUR ENDPOINT URL>"
assert (
    inference_server_url is not "<YOUR ENDPOINT URL>"
), "Please provide your own HF inference endpoint URL!"

llm = HuggingFaceTextGenInference(
    inference_server_url=inference_server_url,
    temperature=0.001,
    max_new_tokens=400,
    do_sample=False,
)
tools = [
    StructuredTool.from_function(calculator),
    StructuredTool.from_function(cylinder_volume),
]
raven_prompt = RavenPromptTemplate(
    template=RAVEN_PROMPT, tools=tools, input_variables=["input"]
)
llm_chain = LLMChain(llm=llm, prompt=raven_prompt)
output_parser = RavenOutputParser()
agent = LLMSingleActionAgent(
    llm_chain=llm_chain,
    output_parser=output_parser,
    stop=["\nReflection:"],
    allowed_tools=tools,
)
agent_chain = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True)

call = agent_chain.run(
    "I have a cake that is about 3 centimenters high and 200 centimeters in radius. How much cake do I have?"
)
call = agent_chain.run("What is 1+10?")
print(exec(call))

scripts/langchain_example.py
nexusflowai/NexusRaven
from typing import Literal

import math

import inspect

from transformers import pipeline


##########################################################
# Step 1: Define the functions you want to articulate. ###
##########################################################


def calculator(
    input_a: float,
    input_b: float,
    operation: Literal["add", "subtract", "multiply", "divide"],
):
    """
    Computes a calculation.

    Args:
    input_a (float) : Required. The first input.
    input_b (float) : Required. The second input.
    operation (string): The operation. Choices include: add to add two numbers, subtract to subtract two numbers, multiply to multiply two numbers, and divide to divide them.
    """
    match operation:
        case "add":
            return input_a + input_b
        case "subtract":
            return input_a - input_b
        case "multiply":
            return input_a * input_b
        case "divide":
            return input_a / input_b


def cylinder_volume(radius, height):
    """
    Calculate the volume of a cylinder.

    Parameters:
    - radius (float): The radius of the base of the cylinder.
    - height (float): The height of the cylinder.

    Returns:
    - float: The volume of the cylinder.
    """
    if radius < 0 or height < 0:
        raise ValueError("Radius and height must be non-negative.")

    volume = math.pi * (radius**2) * height
    return volume


#############################################################
# Step 2: Let's define some utils for building the prompt ###
#############################################################


def format_functions_for_prompt(*functions):
    formatted_functions = []
    for func in functions:
        source_code = inspect.getsource(func)
        docstring = inspect.getdoc(func)
        formatted_functions.append(
            f"OPTION:\n<func_start>{source_code}<func_end>\n<docstring_start>\n{docstring}\n<docstring_end>"
        )
    return "\n".join(formatted_functions)


##############################
# Step 3: Construct Prompt ###
##############################


def construct_prompt(user_query: str):
    formatted_prompt = format_functions_for_prompt(calculator, cylinder_volume)
    formatted_prompt += f"\n\nUser Query: Question: {user_query}\n"

    prompt = (
        "<human>:\n"
        + formatted_prompt
        + "Please pick a function from the above options that best answers the user query and fill in the appropriate arguments.<human_end>"
    )
    return prompt


#######################################
# Step 4: Execute the function call ###
#######################################


def execute_function_call(model_output):
    # Ignore everything after "Reflection" since that is not essential.
    function_call = (
        model_output[0]["generated_text"]
        .strip()
        .split("\n")[1]
        .replace("Initial Answer:", "")
        .strip()
    )

    try:
        return eval(function_call)
    except Exception as e:
        return str(e)


if __name__ == "__main__":
    # Build the model
    text_gen = pipeline(
        "text-generation",
        model="Nexusflow/NexusRaven-13B",
        device="cuda",
    )

    # Comp[ute a Simple equation
    prompt = construct_prompt("What is 1+10?")
    model_output = text_gen(
        prompt, do_sample=False, max_new_tokens=400, return_full_text=False
    )
    result = execute_function_call(model_output)

    print("Model Output:", model_output)
    print("Execution Result:", result)

    prompt = construct_prompt(
        "I have a cake that is about 3 centimenters high and 200 centimeters in diameter. How much cake do I have?"
    )
    model_output = text_gen(
        prompt,
        do_sample=False,
        max_new_tokens=400,
        return_full_text=False,
        stop=["\nReflection:"],
    )
    result = execute_function_call(model_output)

    print("Model Output:", model_output)
    print("Execution Result:", result)

scripts/non_langchain_example.py
nexusflowai/NexusRaven