usrbinenv python codingutf8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license base class for all agents which contains the main api methods args chatprompttemplate str optional pass along your own prompt if you want to override the default template for the chat method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named chatprompttemplate txt in this repo in this case runprompttemplate str optional pass along your own prompt if you want to override the default template for the run method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named runprompttemplate txt in this repo in this case additionaltools tool list of tools or dictionary with tool values optional any additional tools to include on top of the default ones if you pass along a tool with the same name as one of the default tools that default tool will be overridden get all tool currently available to the agent return self toolbox def formatpromptself task chatmodefalse description n joinf name tool description for name tool in self toolbox items if chatmode if self chathistory is none prompt self chatprompttemplate replacealltools description else prompt self chathistory prompt chatmessageprompt replacetask task else prompt self runprompttemplate replacealltools description prompt prompt replaceprompt task return prompt def setstreamself streamer self log streamer def chatself task returncodefalse remotefalse kwargs prompt self formatprompttask chatmodetrue result self generateoneprompt stophuman self chathistory prompt result strip n explanation code cleancodeforchatresult self logfexplanation from the agentnexplanation if code is not none self logfnncode generated by the agentncode if not returncode self lognnresult self cachedtools resolvetoolscode self toolbox remoteremote cachedtoolsself cachedtools self chatstate updatekwargs return evaluatecode self cachedtools self chatstate chatmodetrue else toolcode gettoolcreationcodecode self toolbox remoteremote return ftoolcodencode def preparefornewchatself self chathistory none self chatstate self cachedtools none def runself task returncodefalse remotefalse kwargs prompt self formatprompttask result self generateoneprompt stoptask explanation code cleancodeforrunresult self logfexplanation from the agentnexplanation self logfnncode generated by the agentncode if not returncode self lognnresult self cachedtools resolvetoolscode self toolbox remoteremote cachedtoolsself cachedtools return evaluatecode self cachedtools statekwargs copy else toolcode gettoolcreationcodecode self toolbox remoteremote return ftoolcodencode def generateoneself prompt stop this is the method to implement in your custom agent raise notimplementederror def generatemanyself prompts stop override if you have a way to do batch generation faster than one by one return self generateoneprompt stop for prompt in prompts class openaiagentagent def init self modeltextdavinci003 apikeynone chatprompttemplatenone runprompttemplatenone additionaltoolsnone if not isopenaiavailable raise importerrorusing openaiagent requires openai pip install openai if apikey is none apikey os environ getopenaiapikey none if apikey is none raise valueerror you need an openai key to use openaiagent you can get one here get one here https openai comapi if you have one set it in your env with os environ openaiapikey xxx else openai apikey apikey self model model super init chatprompttemplatechatprompttemplate runprompttemplaterunprompttemplate additionaltoolsadditionaltools def generatemanyself prompts stop if gpt in self model return self chatgenerateprompt stop for prompt in prompts else return self completiongenerateprompts stop def generateoneself prompt stop if gpt in self model return self chatgenerateprompt stop else return self completiongenerateprompt stop0 def chatgenerateself prompt stop result openai chat completions create modelself model messagesrole user content prompt temperature0 stopstop return result choices0 message content def completiongenerateself prompts stop result openai completion create modelself model promptprompts temperature0 stopstop maxtokens200 return answertext for answer in resultchoices class azureopenaiagentagent def init self deploymentid apikeynone resourcenamenone apiversion20221201 ischatmodelnone chatprompttemplatenone runprompttemplatenone additionaltoolsnone if not isopenaiavailable raise importerrorusing openaiagent requires openai pip install openai self deploymentid deploymentid openai apitype azure if apikey is none apikey os environ getazureopenaiapikey none if apikey is none raise valueerror you need an azure openai key to use azureopenaiagent if you have one set it in your env with os environ azureopenaiapikey xxx else openai apikey apikey if resourcename is none resourcename os environ getazureopenairesourcename none if resourcename is none raise valueerror you need a resourcename to use azureopenaiagent if you have one set it in your env with os environ azureopenairesourcename xxx else openai apibase fhttps resourcename openai azure com openai apiversion apiversion if ischatmodel is none ischatmodel gpt in deploymentid lower self ischatmodel ischatmodel super init chatprompttemplatechatprompttemplate runprompttemplaterunprompttemplate additionaltoolsadditionaltools def generatemanyself prompts stop if self ischatmodel return self chatgenerateprompt stop for prompt in prompts else return self completiongenerateprompts stop def generateoneself prompt stop if self ischatmodel return self chatgenerateprompt stop else return self completiongenerateprompt stop0 def chatgenerateself prompt stop result openai chatcompletion create engineself deploymentid messagesrole user content prompt temperature0 stopstop return resultchoices0messagecontent def completiongenerateself prompts stop result openai completion create engineself deploymentid promptprompts temperature0 stopstop maxtokens200 return answertext for answer in resultchoices class hfagentagent def init self urlendpoint tokennone chatprompttemplatenone runprompttemplatenone additionaltoolsnone self urlendpoint urlendpoint if token is none self token fbearer hffolder gettoken elif token startswithbearer or token startswithbasic self token token else self token fbearer token super init chatprompttemplatechatprompttemplate runprompttemplaterunprompttemplate additionaltoolsadditionaltools def generateoneself prompt stop headers ization self token inputs inputs prompt parameters maxnewtokens 200 returnfulltext false stop stop response requests postself urlendpoint jsoninputs headersheaders if response statuscode 429 logger infogetting ratelimited waiting a tiny bit before trying again time sleep1 return self generateoneprompt elif response statuscode 200 raise valueerrorferror response statuscode response json result response json0generatedtext inference api returns the stop sequence for stopseq in stop if result endswithstopseq return result lenstopseq return result class localagentagent def initself model tokenizer chatprompttemplatenone runprompttemplatenone additionaltoolsnone self model model self tokenizer tokenizer super init chatprompttemplatechatprompttemplate runprompttemplaterunprompttemplate additionaltoolsadditionaltools classmethod def frompretrainedcls pretrainedmodelnameorpath kwargs model automodelforcausallm frompretrainedpretrainedmodelnameorpath kwargs tokenizer autotokenizer frompretrainedpretrainedmodelnameorpath kwargs return clsmodel tokenizer property def modeldeviceself if hasattrself model hfdevicemap return listself model hfdevicemap values0 for param in self model parameters return param device def generateoneself prompt stop encodedinputs self tokenizerprompt returntensorspt toself modeldevice srclen encodedinputsinputids shape1 stoppingcriteria stoppingcriterialiststopsequencecriteriastop self tokenizer outputs self model generate encodedinputsinputids maxnewtokens200 stoppingcriteriastoppingcriteria result self tokenizer decodeoutputs0 tolistsrclen inference api returns the stop sequence for stopseq in stop if result endswithstopseq result result lenstopseq return result class stopsequencecriteriastoppingcriteria def initself stopsequences tokenizer if isinstancestopsequences str stopsequences stopsequences self stopsequences stopsequences self tokenizer tokenizer def callself inputids scores kwargs bool decodedoutput self tokenizer decodeinputids tolist0 return anydecodedoutput endswithstopsequence for stopsequence in self stopsequences usr bin env python coding utf 8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license base class for all agents which contains the main api methods args chat_prompt_template str optional pass along your own prompt if you want to override the default template for the chat method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named chat_prompt_template txt in this repo in this case run_prompt_template str optional pass along your own prompt if you want to override the default template for the run method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named run_prompt_template txt in this repo in this case additional_tools tool list of tools or dictionary with tool values optional any additional tools to include on top of the default ones if you pass along a tool with the same name as one of the default tools that default tool will be overridden get all tool currently available to the agent set the function use to stream results which is print by default args streamer callable the function to call when streaming results from the llm sends a new request to the agent in a chat will use the previous ones in its history args task str the task to perform return_code bool optional defaults to false whether to just return code and not evaluate it remote bool optional defaults to false whether or not to use remote tools inference endpoints instead of local ones kwargs additional keyword arguments optional any keyword argument to send to the agent when evaluating the code example py from transformers import hfagent agent hfagent https api inference huggingface co models bigcode starcoder agent chat draw me a picture of rivers and lakes agent chat transform the picture so that there is a rock in there clears the history of prior calls to agent chat sends a request to the agent args task str the task to perform return_code bool optional defaults to false whether to just return code and not evaluate it remote bool optional defaults to false whether or not to use remote tools inference endpoints instead of local ones kwargs additional keyword arguments optional any keyword argument to send to the agent when evaluating the code example py from transformers import hfagent agent hfagent https api inference huggingface co models bigcode starcoder agent run draw me a picture of rivers and lakes this is the method to implement in your custom agent override if you have a way to do batch generation faster than one by one agent that uses the openai api to generate code tip warning true the openai models are used in generation mode so even for the chat api it s better to use models like text davinci 003 over the chat gpt variant proper support for chat gpt models will come in a next version tip args model str optional defaults to text davinci 003 the name of the openai model to use api_key str optional the api key to use if unset will look for the environment variable openai_api_key chat_prompt_template str optional pass along your own prompt if you want to override the default template for the chat method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named chat_prompt_template txt in this repo in this case run_prompt_template str optional pass along your own prompt if you want to override the default template for the run method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named run_prompt_template txt in this repo in this case additional_tools tool list of tools or dictionary with tool values optional any additional tools to include on top of the default ones if you pass along a tool with the same name as one of the default tools that default tool will be overridden example py from transformers import openaiagent agent openaiagent model text davinci 003 api_key xxx agent run is the following text in spanish positive or negative text este es un api muy agradable agent that uses azure openai to generate code see the official documentation https learn microsoft com en us azure cognitive services openai to learn how to deploy an openai model on azure tip warning true the openai models are used in generation mode so even for the chat api it s better to use models like text davinci 003 over the chat gpt variant proper support for chat gpt models will come in a next version tip args deployment_id str the name of the deployed azure openai model to use api_key str optional the api key to use if unset will look for the environment variable azure_openai_api_key resource_name str optional the name of your azure openai resource if unset will look for the environment variable azure_openai_resource_name api_version str optional default to 2022 12 01 the api version to use for this agent is_chat_mode bool optional whether you are using a completion model or a chat model see note above chat models won t be as efficient will default to gpt being in the deployment_id or not chat_prompt_template str optional pass along your own prompt if you want to override the default template for the chat method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named chat_prompt_template txt in this repo in this case run_prompt_template str optional pass along your own prompt if you want to override the default template for the run method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named run_prompt_template txt in this repo in this case additional_tools tool list of tools or dictionary with tool values optional any additional tools to include on top of the default ones if you pass along a tool with the same name as one of the default tools that default tool will be overridden example py from transformers import azureopenaiagent agent azureaiagent deployment_id davinci 003 api_key xxx resource_name yyy agent run is the following text in spanish positive or negative text este es un api muy agradable agent that uses an inference endpoint to generate code args url_endpoint str the name of the url endpoint to use token str optional the token to use as http bearer ization for remote files if unset will use the token generated when running huggingface cli login stored in huggingface chat_prompt_template str optional pass along your own prompt if you want to override the default template for the chat method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named chat_prompt_template txt in this repo in this case run_prompt_template str optional pass along your own prompt if you want to override the default template for the run method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named run_prompt_template txt in this repo in this case additional_tools tool list of tools or dictionary with tool values optional any additional tools to include on top of the default ones if you pass along a tool with the same name as one of the default tools that default tool will be overridden example py from transformers import hfagent agent hfagent https api inference huggingface co models bigcode starcoder agent run is the following text in spanish positive or negative text este es un api muy agradable inference api returns the stop sequence agent that uses a local model and tokenizer to generate code args model pretrainedmodel the model to use for the agent tokenizer pretrainedtokenizer the tokenizer to use for the agent chat_prompt_template str optional pass along your own prompt if you want to override the default template for the chat method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named chat_prompt_template txt in this repo in this case run_prompt_template str optional pass along your own prompt if you want to override the default template for the run method can be the actual prompt template or a repo id on the hugging face hub the prompt should be in a file named run_prompt_template txt in this repo in this case additional_tools tool list of tools or dictionary with tool values optional any additional tools to include on top of the default ones if you pass along a tool with the same name as one of the default tools that default tool will be overridden example py import torch from transformers import automodelforcausallm autotokenizer localagent checkpoint bigcode starcoder model automodelforcausallm from_pretrained checkpoint device_map auto torch_dtype torch bfloat16 tokenizer autotokenizer from_pretrained checkpoint agent localagent model tokenizer agent run draw me a picture of rivers and lakes convenience method to build a localagent from a pretrained checkpoint args pretrained_model_name_or_path str or os pathlike the name of a repo on the hub or a local path to a folder containing both model and tokenizer kwargs dict str any optional keyword arguments passed along to pretrainedmodel from_pretrained example py import torch from transformers import localagent agent localagent from_pretrained bigcode starcoder device_map auto torch_dtype torch bfloat16 agent run draw me a picture of rivers and lakes inference api returns the stop sequence this class can be used to stop generation whenever a sequence of tokens is encountered args stop_sequences str or list str the sequence or list of sequences on which to stop execution tokenizer the tokenizer used to decode the model outputs
import importlib.util
import json
import os
import time
from dataclasses import dataclass
from typing import Dict

import requests
from huggingface_hub import HfFolder, hf_hub_download, list_spaces

from ..models.auto import AutoTokenizer
from ..utils import is_offline_mode, is_openai_available, is_torch_available, logging
from .base import TASK_MAPPING, TOOL_CONFIG_FILE, Tool, load_tool, supports_remote
from .prompts import CHAT_MESSAGE_PROMPT, download_prompt
from .python_interpreter import evaluate


logger = logging.get_logger(__name__)


if is_openai_available():
    import openai

if is_torch_available():
    from ..generation import StoppingCriteria, StoppingCriteriaList
    from ..models.auto import AutoModelForCausalLM
else:
    StoppingCriteria = object

_tools_are_initialized = False


BASE_PYTHON_TOOLS = {
    "print": print,
    "range": range,
    "float": float,
    "int": int,
    "bool": bool,
    "str": str,
}


@dataclass
class PreTool:
    task: str
    description: str
    repo_id: str


HUGGINGFACE_DEFAULT_TOOLS = {}


HUGGINGFACE_DEFAULT_TOOLS_FROM_HUB = [
    "image-transformation",
    "text-download",
    "text-to-image",
    "text-to-video",
]


def get_remote_tools(organization="huggingface-tools"):
    if is_offline_mode():
        logger.info("You are in offline mode, so remote tools are not available.")
        return {}

    spaces = list_spaces(author=organization)
    tools = {}
    for space_info in spaces:
        repo_id = space_info.id
        resolved_config_file = hf_hub_download(repo_id, TOOL_CONFIG_FILE, repo_type="space")
        with open(resolved_config_file, encoding="utf-8") as reader:
            config = json.load(reader)

        task = repo_id.split("/")[-1]
        tools[config["name"]] = PreTool(task=task, description=config["description"], repo_id=repo_id)

    return tools


def _setup_default_tools():
    global HUGGINGFACE_DEFAULT_TOOLS
    global _tools_are_initialized

    if _tools_are_initialized:
        return

    main_module = importlib.import_module("transformers")
    tools_module = main_module.tools

    remote_tools = get_remote_tools()
    for task_name, tool_class_name in TASK_MAPPING.items():
        tool_class = getattr(tools_module, tool_class_name)
        description = tool_class.description
        HUGGINGFACE_DEFAULT_TOOLS[tool_class.name] = PreTool(task=task_name, description=description, repo_id=None)

    if not is_offline_mode():
        for task_name in HUGGINGFACE_DEFAULT_TOOLS_FROM_HUB:
            found = False
            for tool_name, tool in remote_tools.items():
                if tool.task == task_name:
                    HUGGINGFACE_DEFAULT_TOOLS[tool_name] = tool
                    found = True
                    break

            if not found:
                raise ValueError(f"{task_name} is not implemented on the Hub.")

    _tools_are_initialized = True


def resolve_tools(code, toolbox, remote=False, cached_tools=None):
    if cached_tools is None:
        resolved_tools = BASE_PYTHON_TOOLS.copy()
    else:
        resolved_tools = cached_tools
    for name, tool in toolbox.items():
        if name not in code or name in resolved_tools:
            continue

        if isinstance(tool, Tool):
            resolved_tools[name] = tool
        else:
            task_or_repo_id = tool.task if tool.repo_id is None else tool.repo_id
            _remote = remote and supports_remote(task_or_repo_id)
            resolved_tools[name] = load_tool(task_or_repo_id, remote=_remote)

    return resolved_tools


def get_tool_creation_code(code, toolbox, remote=False):
    code_lines = ["from transformers import load_tool", ""]
    for name, tool in toolbox.items():
        if name not in code or isinstance(tool, Tool):
            continue

        task_or_repo_id = tool.task if tool.repo_id is None else tool.repo_id
        line = f'{name} = load_tool("{task_or_repo_id}"'
        if remote:
            line += ", remote=True"
        line += ")"
        code_lines.append(line)

    return "\n".join(code_lines) + "\n"


def clean_code_for_chat(result):
    lines = result.split("\n")
    idx = 0
    while idx < len(lines) and not lines[idx].lstrip().startswith("```"):
        idx += 1
    explanation = "\n".join(lines[:idx]).strip()
    if idx == len(lines):
        return explanation, None

    idx += 1
    start_idx = idx
    while not lines[idx].lstrip().startswith("```"):
        idx += 1
    code = "\n".join(lines[start_idx:idx]).strip()

    return explanation, code


def clean_code_for_run(result):
    result = f"I will use the following {result}"
    explanation, code = result.split("Answer:")
    explanation = explanation.strip()
    code = code.strip()

    code_lines = code.split("\n")
    if code_lines[0] in ["```", "```py", "```python"]:
        code_lines = code_lines[1:]
    if code_lines[-1] == "```":
        code_lines = code_lines[:-1]
    code = "\n".join(code_lines)

    return explanation, code


class Agent:
    

    def __init__(self, chat_prompt_template=None, run_prompt_template=None, additional_tools=None):
        _setup_default_tools()

        agent_name = self.__class__.__name__
        self.chat_prompt_template = download_prompt(chat_prompt_template, agent_name, mode="chat")
        self.run_prompt_template = download_prompt(run_prompt_template, agent_name, mode="run")
        self._toolbox = HUGGINGFACE_DEFAULT_TOOLS.copy()
        self.log = print
        if additional_tools is not None:
            if isinstance(additional_tools, (list, tuple)):
                additional_tools = {t.name: t for t in additional_tools}
            elif not isinstance(additional_tools, dict):
                additional_tools = {additional_tools.name: additional_tools}

            replacements = {name: tool for name, tool in additional_tools.items() if name in HUGGINGFACE_DEFAULT_TOOLS}
            self._toolbox.update(additional_tools)
            if len(replacements) > 1:
                names = "\n".join([f"- {n}: {t}" for n, t in replacements.items()])
                logger.warning(
                    f"The following tools have been replaced by the ones provided in `additional_tools`:\n{names}."
                )
            elif len(replacements) == 1:
                name = list(replacements.keys())[0]
                logger.warning(f"{name} has been replaced by {replacements[name]} as provided in `additional_tools`.")

        self.prepare_for_new_chat()

    @property
    def toolbox(self) -> Dict[str, Tool]:
        
        return self._toolbox

    def format_prompt(self, task, chat_mode=False):
        description = "\n".join([f"- {name}: {tool.description}" for name, tool in self.toolbox.items()])
        if chat_mode:
            if self.chat_history is None:
                prompt = self.chat_prompt_template.replace("<<all_tools>>", description)
            else:
                prompt = self.chat_history
            prompt += CHAT_MESSAGE_PROMPT.replace("<<task>>", task)
        else:
            prompt = self.run_prompt_template.replace("<<all_tools>>", description)
            prompt = prompt.replace("<<prompt>>", task)
        return prompt

    def set_stream(self, streamer):
        
        self.log = streamer

    def chat(self, task, *, return_code=False, remote=False, **kwargs):
        
        prompt = self.format_prompt(task, chat_mode=True)
        result = self.generate_one(prompt, stop=["Human:", "====="])
        self.chat_history = prompt + result.strip() + "\n"
        explanation, code = clean_code_for_chat(result)

        self.log(f"==Explanation from the agent==\n{explanation}")

        if code is not None:
            self.log(f"\n\n==Code generated by the agent==\n{code}")
            if not return_code:
                self.log("\n\n==Result==")
                self.cached_tools = resolve_tools(code, self.toolbox, remote=remote, cached_tools=self.cached_tools)
                self.chat_state.update(kwargs)
                return evaluate(code, self.cached_tools, self.chat_state, chat_mode=True)
            else:
                tool_code = get_tool_creation_code(code, self.toolbox, remote=remote)
                return f"{tool_code}\n{code}"

    def prepare_for_new_chat(self):
        
        self.chat_history = None
        self.chat_state = {}
        self.cached_tools = None

    def run(self, task, *, return_code=False, remote=False, **kwargs):
        
        prompt = self.format_prompt(task)
        result = self.generate_one(prompt, stop=["Task:"])
        explanation, code = clean_code_for_run(result)

        self.log(f"==Explanation from the agent==\n{explanation}")

        self.log(f"\n\n==Code generated by the agent==\n{code}")
        if not return_code:
            self.log("\n\n==Result==")
            self.cached_tools = resolve_tools(code, self.toolbox, remote=remote, cached_tools=self.cached_tools)
            return evaluate(code, self.cached_tools, state=kwargs.copy())
        else:
            tool_code = get_tool_creation_code(code, self.toolbox, remote=remote)
            return f"{tool_code}\n{code}"

    def generate_one(self, prompt, stop):
        
        raise NotImplementedError

    def generate_many(self, prompts, stop):
        
        return [self.generate_one(prompt, stop) for prompt in prompts]


class OpenAiAgent(Agent):
    

    def __init__(
        self,
        model="text-davinci-003",
        api_key=None,
        chat_prompt_template=None,
        run_prompt_template=None,
        additional_tools=None,
    ):
        if not is_openai_available():
            raise ImportError("Using `OpenAiAgent` requires `openai`: `pip install openai`.")

        if api_key is None:
            api_key = os.environ.get("OPENAI_API_KEY", None)
        if api_key is None:
            raise ValueError(
                "You need an openai key to use `OpenAIAgent`. You can get one here: Get one here "
                "https://openai.com/api/`. If you have one, set it in your env with `os.environ['OPENAI_API_KEY'] = "
                "xxx."
            )
        else:
            openai.api_key = api_key
        self.model = model
        super().__init__(
            chat_prompt_template=chat_prompt_template,
            run_prompt_template=run_prompt_template,
            additional_tools=additional_tools,
        )

    def generate_many(self, prompts, stop):
        if "gpt" in self.model:
            return [self._chat_generate(prompt, stop) for prompt in prompts]
        else:
            return self._completion_generate(prompts, stop)

    def generate_one(self, prompt, stop):
        if "gpt" in self.model:
            return self._chat_generate(prompt, stop)
        else:
            return self._completion_generate([prompt], stop)[0]

    def _chat_generate(self, prompt, stop):
        result = openai.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            stop=stop,
        )
        return result.choices[0].message.content

    def _completion_generate(self, prompts, stop):
        result = openai.Completion.create(
            model=self.model,
            prompt=prompts,
            temperature=0,
            stop=stop,
            max_tokens=200,
        )
        return [answer["text"] for answer in result["choices"]]


class AzureOpenAiAgent(Agent):
    

    def __init__(
        self,
        deployment_id,
        api_key=None,
        resource_name=None,
        api_version="2022-12-01",
        is_chat_model=None,
        chat_prompt_template=None,
        run_prompt_template=None,
        additional_tools=None,
    ):
        if not is_openai_available():
            raise ImportError("Using `OpenAiAgent` requires `openai`: `pip install openai`.")

        self.deployment_id = deployment_id
        openai.api_type = "azure"
        if api_key is None:
            api_key = os.environ.get("AZURE_OPENAI_API_KEY", None)
        if api_key is None:
            raise ValueError(
                "You need an Azure openAI key to use `AzureOpenAIAgent`. If you have one, set it in your env with "
                "`os.environ['AZURE_OPENAI_API_KEY'] = xxx."
            )
        else:
            openai.api_key = api_key
        if resource_name is None:
            resource_name = os.environ.get("AZURE_OPENAI_RESOURCE_NAME", None)
        if resource_name is None:
            raise ValueError(
                "You need a resource_name to use `AzureOpenAIAgent`. If you have one, set it in your env with "
                "`os.environ['AZURE_OPENAI_RESOURCE_NAME'] = xxx."
            )
        else:
            openai.api_base = f"https://{resource_name}.openai.azure.com"
        openai.api_version = api_version

        if is_chat_model is None:
            is_chat_model = "gpt" in deployment_id.lower()
        self.is_chat_model = is_chat_model

        super().__init__(
            chat_prompt_template=chat_prompt_template,
            run_prompt_template=run_prompt_template,
            additional_tools=additional_tools,
        )

    def generate_many(self, prompts, stop):
        if self.is_chat_model:
            return [self._chat_generate(prompt, stop) for prompt in prompts]
        else:
            return self._completion_generate(prompts, stop)

    def generate_one(self, prompt, stop):
        if self.is_chat_model:
            return self._chat_generate(prompt, stop)
        else:
            return self._completion_generate([prompt], stop)[0]

    def _chat_generate(self, prompt, stop):
        result = openai.ChatCompletion.create(
            engine=self.deployment_id,
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            stop=stop,
        )
        return result["choices"][0]["message"]["content"]

    def _completion_generate(self, prompts, stop):
        result = openai.Completion.create(
            engine=self.deployment_id,
            prompt=prompts,
            temperature=0,
            stop=stop,
            max_tokens=200,
        )
        return [answer["text"] for answer in result["choices"]]


class HfAgent(Agent):
    

    def __init__(
        self, url_endpoint, token=None, chat_prompt_template=None, run_prompt_template=None, additional_tools=None
    ):
        self.url_endpoint = url_endpoint
        if token is None:
            self.token = f"Bearer {HfFolder().get_token()}"
        elif token.startswith("Bearer") or token.startswith("Basic"):
            self.token = token
        else:
            self.token = f"Bearer {token}"
        super().__init__(
            chat_prompt_template=chat_prompt_template,
            run_prompt_template=run_prompt_template,
            additional_tools=additional_tools,
        )

    def generate_one(self, prompt, stop):
        headers = {"Authorization": self.token}
        inputs = {
            "inputs": prompt,
            "parameters": {"max_new_tokens": 200, "return_full_text": False, "stop": stop},
        }

        response = requests.post(self.url_endpoint, json=inputs, headers=headers)
        if response.status_code == 429:
            logger.info("Getting rate-limited, waiting a tiny bit before trying again.")
            time.sleep(1)
            return self._generate_one(prompt)
        elif response.status_code != 200:
            raise ValueError(f"Error {response.status_code}: {response.json()}")

        result = response.json()[0]["generated_text"]
        
        for stop_seq in stop:
            if result.endswith(stop_seq):
                return result[: -len(stop_seq)]
        return result


class LocalAgent(Agent):
    

    def __init__(self, model, tokenizer, chat_prompt_template=None, run_prompt_template=None, additional_tools=None):
        self.model = model
        self.tokenizer = tokenizer
        super().__init__(
            chat_prompt_template=chat_prompt_template,
            run_prompt_template=run_prompt_template,
            additional_tools=additional_tools,
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        
        model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path, **kwargs)
        tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, **kwargs)
        return cls(model, tokenizer)

    @property
    def _model_device(self):
        if hasattr(self.model, "hf_device_map"):
            return list(self.model.hf_device_map.values())[0]
        for param in self.model.parameters():
            return param.device

    def generate_one(self, prompt, stop):
        encoded_inputs = self.tokenizer(prompt, return_tensors="pt").to(self._model_device)
        src_len = encoded_inputs["input_ids"].shape[1]
        stopping_criteria = StoppingCriteriaList([StopSequenceCriteria(stop, self.tokenizer)])
        outputs = self.model.generate(
            encoded_inputs["input_ids"], max_new_tokens=200, stopping_criteria=stopping_criteria
        )

        result = self.tokenizer.decode(outputs[0].tolist()[src_len:])
        
        for stop_seq in stop:
            if result.endswith(stop_seq):
                result = result[: -len(stop_seq)]
        return result


class StopSequenceCriteria(StoppingCriteria):
    

    def __init__(self, stop_sequences, tokenizer):
        if isinstance(stop_sequences, str):
            stop_sequences = [stop_sequences]
        self.stop_sequences = stop_sequences
        self.tokenizer = tokenizer

    def __call__(self, input_ids, scores, **kwargs) -> bool:
        decoded_output = self.tokenizer.decode(input_ids.tolist()[0])
        return any(decoded_output.endswith(stop_sequence) for stop_sequence in self.stop_sequences)