CoolFace
Apppublic

dkolarova/bpm-agent

sourceHugging Faceupdated 2y agoView on Hugging Face
0likes
models.py851 linesDownload Raw Back to root
1import json2import logging3import os4import random5from copy import deepcopy6from dataclasses import asdict, dataclass7from enum import Enum8from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union9 10from huggingface_hub import InferenceClient11from huggingface_hub.utils import is_torch_available12from PIL import Image13 14from smolagents.tools import Tool15 16 17 18if TYPE_CHECKING:19    from transformers import StoppingCriteriaList20 21logger = logging.getLogger(__name__)22 23DEFAULT_JSONAGENT_REGEX_GRAMMAR = {24    "type": "regex",25    "value": 'Thought: .+?\\nAction:\\n\\{\\n\\s{4}"action":\\s"[^"\\n]+",\\n\\s{4}"action_input":\\s"[^"\\n]+"\\n\\}\\n<end_code>',26}27 28DEFAULT_CODEAGENT_REGEX_GRAMMAR = {29    "type": "regex",30    "value": "Thought: .+?\\nCode:\\n```(?:py|python)?\\n(?:.|\\s)+?\\n```<end_code>",31}32 33 34def _is_package_available(package_name: str) -> bool:35    try:36        importlib.metadata.version(package_name)37        return True38    except importlib.metadata.PackageNotFoundError:39        return False40 41def encode_image_base64(image):42    buffered = BytesIO()43    image.save(buffered, format="PNG")44    return base64.b64encode(buffered.getvalue()).decode("utf-8")45 46 47def make_image_url(base64_image):48    return f"data:image/png;base64,{base64_image}"49 50def get_dict_from_nested_dataclasses(obj, ignore_key=None):51    def convert(obj):52        if hasattr(obj, "__dataclass_fields__"):53            return {k: convert(v) for k, v in asdict(obj).items() if k != ignore_key}54        return obj55 56    return convert(obj)57 58 59@dataclass60class ChatMessageToolCallDefinition:61    arguments: Any62    name: str63    description: Optional[str] = None64 65    @classmethod66    def from_hf_api(cls, tool_call_definition) -> "ChatMessageToolCallDefinition":67        return cls(68            arguments=tool_call_definition.arguments,69            name=tool_call_definition.name,70            description=tool_call_definition.description,71        )72 73 74@dataclass75class ChatMessageToolCall:76    function: ChatMessageToolCallDefinition77    id: str78    type: str79 80    @classmethod81    def from_hf_api(cls, tool_call) -> "ChatMessageToolCall":82        return cls(83            function=ChatMessageToolCallDefinition.from_hf_api(tool_call.function),84            id=tool_call.id,85            type=tool_call.type,86        )87 88 89@dataclass90class ChatMessage:91    role: str92    content: Optional[str] = None93    tool_calls: Optional[List[ChatMessageToolCall]] = None94    raw: Optional[Any] = None  # Stores the raw output from the API95 96    def model_dump_json(self):97        return json.dumps(get_dict_from_nested_dataclasses(self, ignore_key="raw"))98 99    @classmethod100    def from_hf_api(cls, message, raw) -> "ChatMessage":101        tool_calls = None102        if getattr(message, "tool_calls", None) is not None:103            tool_calls = [ChatMessageToolCall.from_hf_api(tool_call) for tool_call in message.tool_calls]104        return cls(role=message.role, content=message.content, tool_calls=tool_calls, raw=raw)105 106    @classmethod107    def from_dict(cls, data: dict) -> "ChatMessage":108        if data.get("tool_calls"):109            tool_calls = [110                ChatMessageToolCall(111                    function=ChatMessageToolCallDefinition(**tc["function"]), id=tc["id"], type=tc["type"]112                )113                for tc in data["tool_calls"]114            ]115            data["tool_calls"] = tool_calls116        return cls(**data)117 118    def dict(self):119        return json.dumps(get_dict_from_nested_dataclasses(self))120 121 122def parse_json_if_needed(arguments: Union[str, dict]) -> Union[str, dict]:123    if isinstance(arguments, dict):124        return arguments125    else:126        try:127            return json.loads(arguments)128        except Exception:129            return arguments130 131 132def parse_tool_args_if_needed(message: ChatMessage) -> ChatMessage:133    for tool_call in message.tool_calls:134        tool_call.function.arguments = parse_json_if_needed(tool_call.function.arguments)135    return message136 137 138class MessageRole(str, Enum):139    USER = "user"140    ASSISTANT = "assistant"141    SYSTEM = "system"142    TOOL_CALL = "tool-call"143    TOOL_RESPONSE = "tool-response"144 145    @classmethod146    def roles(cls):147        return [r.value for r in cls]148 149 150tool_role_conversions = {151    MessageRole.TOOL_CALL: MessageRole.ASSISTANT,152    MessageRole.TOOL_RESPONSE: MessageRole.USER,153}154 155 156def get_tool_json_schema(tool: Tool) -> Dict:157    properties = deepcopy(tool.inputs)158    required = []159    for key, value in properties.items():160        if value["type"] == "any":161            value["type"] = "string"162        if not ("nullable" in value and value["nullable"]):163            required.append(key)164    return {165        "type": "function",166        "function": {167            "name": tool.name,168            "description": tool.description,169            "parameters": {170                "type": "object",171                "properties": properties,172                "required": required,173            },174        },175    }176 177 178def remove_stop_sequences(content: str, stop_sequences: List[str]) -> str:179    for stop_seq in stop_sequences:180        if content[-len(stop_seq) :] == stop_seq:181            content = content[: -len(stop_seq)]182    return content183 184 185def get_clean_message_list(186    message_list: List[Dict[str, str]],187    role_conversions: Dict[MessageRole, MessageRole] = {},188    convert_images_to_image_urls: bool = False,189    flatten_messages_as_text: bool = False,190) -> List[Dict[str, str]]:191    """192    Subsequent messages with the same role will be concatenated to a single message.193    output_message_list is a list of messages that will be used to generate the final message that is chat template compatible with transformers LLM chat template.194 195    Args:196        message_list (`list[dict[str, str]]`): List of chat messages.197        role_conversions (`dict[MessageRole, MessageRole]`, *optional* ): Mapping to convert roles.198        convert_images_to_image_urls (`bool`, default `False`): Whether to convert images to image URLs.199        flatten_messages_as_text (`bool`, default `False`): Whether to flatten messages as text.200    """201    output_message_list = []202    message_list = deepcopy(message_list)  # Avoid modifying the original list203    for message in message_list:204        role = message["role"]205        if role not in MessageRole.roles():206            raise ValueError(f"Incorrect role {role}, only {MessageRole.roles()} are supported for now.")207 208        if role in role_conversions:209            message["role"] = role_conversions[role]210        # encode images if needed211        if isinstance(message["content"], list):212            for element in message["content"]:213                if element["type"] == "image":214                    assert not flatten_messages_as_text, f"Cannot use images with {flatten_messages_as_text=}"215                    if convert_images_to_image_urls:216                        element.update(217                            {218                                "type": "image_url",219                                "image_url": {"url": make_image_url(encode_image_base64(element.pop("image")))},220                            }221                        )222                    else:223                        element["image"] = encode_image_base64(element["image"])224 225        if len(output_message_list) > 0 and message["role"] == output_message_list[-1]["role"]:226            assert isinstance(message["content"], list), "Error: wrong content:" + str(message["content"])227            if flatten_messages_as_text:228                output_message_list[-1]["content"] += message["content"][0]["text"]229            else:230                output_message_list[-1]["content"] += message["content"]231        else:232            if flatten_messages_as_text:233                content = message["content"][0]["text"]234            else:235                content = message["content"]236            output_message_list.append({"role": message["role"], "content": content})237    return output_message_list238 239 240class Model:241    def __init__(self, **kwargs):242        self.last_input_token_count = None243        self.last_output_token_count = None244        self.kwargs = kwargs245 246    def _prepare_completion_kwargs(247        self,248        messages: List[Dict[str, str]],249        stop_sequences: Optional[List[str]] = None,250        grammar: Optional[str] = None,251        tools_to_call_from: Optional[List[Tool]] = None,252        custom_role_conversions: Optional[Dict[str, str]] = None,253        convert_images_to_image_urls: bool = False,254        flatten_messages_as_text: bool = False,255        **kwargs,256    ) -> Dict:257        """258        Prepare parameters required for model invocation, handling parameter priorities.259 260        Parameter priority from high to low:261        1. Explicitly passed kwargs262        2. Specific parameters (stop_sequences, grammar, etc.)263        3. Default values in self.kwargs264        """265        # Clean and standardize the message list266        messages = get_clean_message_list(267            messages,268            role_conversions=custom_role_conversions or tool_role_conversions,269            convert_images_to_image_urls=convert_images_to_image_urls,270            flatten_messages_as_text=flatten_messages_as_text,271        )272 273        # Use self.kwargs as the base configuration274        completion_kwargs = {275            **self.kwargs,276            "messages": messages,277        }278 279        # Handle specific parameters280        if stop_sequences is not None:281            completion_kwargs["stop"] = stop_sequences282        if grammar is not None:283            completion_kwargs["grammar"] = grammar284 285        # Handle tools parameter286        if tools_to_call_from:287            completion_kwargs.update(288                {289                    "tools": [get_tool_json_schema(tool) for tool in tools_to_call_from],290                    "tool_choice": "required",291                }292            )293 294        # Finally, use the passed-in kwargs to override all settings295        completion_kwargs.update(kwargs)296 297        return completion_kwargs298 299    def get_token_counts(self) -> Dict[str, int]:300        return {301            "input_token_count": self.last_input_token_count,302            "output_token_count": self.last_output_token_count,303        }304 305    def __call__(306        self,307        messages: List[Dict[str, str]],308        stop_sequences: Optional[List[str]] = None,309        grammar: Optional[str] = None,310        tools_to_call_from: Optional[List[Tool]] = None,311        **kwargs,312    ) -> ChatMessage:313        """Process the input messages and return the model's response.314 315        Parameters:316            messages (`List[Dict[str, str]]`):317                A list of message dictionaries to be processed. Each dictionary should have the structure `{"role": "user/system", "content": "message content"}`.318            stop_sequences (`List[str]`, *optional*):319                A list of strings that will stop the generation if encountered in the model's output.320            grammar (`str`, *optional*):321                The grammar or formatting structure to use in the model's response.322            tools_to_call_from (`List[Tool]`, *optional*):323                A list of tools that the model can use to generate responses.324            **kwargs:325                Additional keyword arguments to be passed to the underlying model.326 327        Returns:328            `ChatMessage`: A chat message object containing the model's response.329        """330        pass  # To be implemented in child classes!331 332 333class HfApiModel(Model):334    """A class to interact with Hugging Face's Inference API for language model interaction.335 336    This model allows you to communicate with Hugging Face's models using the Inference API. It can be used in both serverless mode or with a dedicated endpoint, supporting features like stop sequences and grammar customization.337 338    Parameters:339        model_id (`str`, *optional*, defaults to `"Qwen/Qwen2.5-Coder-32B-Instruct"`):340            The Hugging Face model ID to be used for inference. This can be a path or model identifier from the Hugging Face model hub.341        provider (`str`, *optional*):342            Name of the provider to use for inference. Can be `"replicate"`, `"together"`, `"fal-ai"`, `"sambanova"` or `"hf-inference"`.343            defaults to hf-inference (HF Inference API).344        token (`str`, *optional*):345            Token used by the Hugging Face API for authentication. This token need to be authorized 'Make calls to the serverless Inference API'.346            If the model is gated (like Llama-3 models), the token also needs 'Read access to contents of all public gated repos you can access'.347            If not provided, the class will try to use environment variable 'HF_TOKEN', else use the token stored in the Hugging Face CLI configuration.348        timeout (`int`, *optional*, defaults to 120):349            Timeout for the API request, in seconds.350        custom_role_conversions (`dict[str, str]`, *optional*):351            Custom role conversion mapping to convert message roles in others.352            Useful for specific models that do not support specific message roles like "system".353        **kwargs:354            Additional keyword arguments to pass to the Hugging Face API.355 356    Raises:357        ValueError:358            If the model name is not provided.359 360    Example:361    ```python362    >>> engine = HfApiModel(363    ...     model_id="Qwen/Qwen2.5-Coder-32B-Instruct",364    ...     token="your_hf_token_here",365    ...     max_tokens=5000,366    ... )367    >>> messages = [{"role": "user", "content": "Explain quantum mechanics in simple terms."}]368    >>> response = engine(messages, stop_sequences=["END"])369    >>> print(response)370    "Quantum mechanics is the branch of physics that studies..."371    ```372    """373 374    def __init__(375        self,376        model_id: str = "Qwen/Qwen2.5-Coder-32B-Instruct",377        provider: Optional[str] = None,378        token: Optional[str] = None,379        timeout: Optional[int] = 120,380        custom_role_conversions: Optional[Dict[str, str]] = None,381        **kwargs,382    ):383        super().__init__(**kwargs)384        self.model_id = model_id385        self.provider = provider386        if token is None:387            token = os.getenv("HF_TOKEN")388        self.client = InferenceClient(self.model_id, provider=provider, token=token, timeout=timeout)389        self.custom_role_conversions = custom_role_conversions390 391    def __call__(392        self,393        messages: List[Dict[str, str]],394        stop_sequences: Optional[List[str]] = None,395        grammar: Optional[str] = None,396        tools_to_call_from: Optional[List[Tool]] = None,397        **kwargs,398    ) -> ChatMessage:399        completion_kwargs = self._prepare_completion_kwargs(400            messages=messages,401            stop_sequences=stop_sequences,402            grammar=grammar,403            tools_to_call_from=tools_to_call_from,404            convert_images_to_image_urls=True,405            custom_role_conversions=self.custom_role_conversions,406            **kwargs,407        )408        response = self.client.chat_completion(**completion_kwargs)409        print(completion_kwargs)410        print('===================================================')411        print(response)412        print('===================================================')413        logger.debug(response)414 415        self.last_input_token_count = response.usage.prompt_tokens416        self.last_output_token_count = response.usage.completion_tokens417        message = ChatMessage.from_hf_api(response.choices[0].message, raw=response)418        if tools_to_call_from is not None:419            return parse_tool_args_if_needed(message)420        return message421 422 423class TransformersModel(Model):424    """A class that uses Hugging Face's Transformers library for language model interaction.425 426    This model allows you to load and use Hugging Face's models locally using the Transformers library. It supports features like stop sequences and grammar customization.427 428    > [!TIP]429    > You must have `transformers` and `torch` installed on your machine. Please run `pip install smolagents[transformers]` if it's not the case.430 431    Parameters:432        model_id (`str`, *optional*, defaults to `"Qwen/Qwen2.5-Coder-32B-Instruct"`):433            The Hugging Face model ID to be used for inference. This can be a path or model identifier from the Hugging Face model hub.434        device_map (`str`, *optional*):435            The device_map to initialize your model with.436        torch_dtype (`str`, *optional*):437            The torch_dtype to initialize your model with.438        trust_remote_code (bool, default `False`):439            Some models on the Hub require running remote code: for this model, you would have to set this flag to True.440        kwargs (dict, *optional*):441            Any additional keyword arguments that you want to use in model.generate(), for instance `max_new_tokens` or `device`.442        **kwargs:443            Additional keyword arguments to pass to `model.generate()`, for instance `max_new_tokens` or `device`.444    Raises:445        ValueError:446            If the model name is not provided.447 448    Example:449    ```python450    >>> engine = TransformersModel(451    ...     model_id="Qwen/Qwen2.5-Coder-32B-Instruct",452    ...     device="cuda",453    ...     max_new_tokens=5000,454    ... )455    >>> messages = [{"role": "user", "content": "Explain quantum mechanics in simple terms."}]456    >>> response = engine(messages, stop_sequences=["END"])457    >>> print(response)458    "Quantum mechanics is the branch of physics that studies..."459    ```460    """461 462    def __init__(463        self,464        model_id: Optional[str] = None,465        device_map: Optional[str] = None,466        torch_dtype: Optional[str] = None,467        trust_remote_code: bool = False,468        **kwargs,469    ):470        super().__init__(**kwargs)471        if not is_torch_available() or not _is_package_available("transformers"):472            raise ModuleNotFoundError(473                "Please install 'transformers' extra to use 'TransformersModel': `pip install 'smolagents[transformers]'`"474            )475        import torch476        from transformers import AutoModelForCausalLM, AutoModelForImageTextToText, AutoProcessor, AutoTokenizer477 478        default_model_id = "HuggingFaceTB/SmolLM2-1.7B-Instruct"479        if model_id is None:480            model_id = default_model_id481            logger.warning(f"`model_id`not provided, using this default tokenizer for token counts: '{model_id}'")482        self.model_id = model_id483        self.kwargs = kwargs484        if device_map is None:485            device_map = "cuda" if torch.cuda.is_available() else "cpu"486        logger.info(f"Using device: {device_map}")487        self._is_vlm = False488        try:489            self.model = AutoModelForCausalLM.from_pretrained(490                model_id,491                device_map=device_map,492                torch_dtype=torch_dtype,493                trust_remote_code=trust_remote_code,494            )495            self.tokenizer = AutoTokenizer.from_pretrained(model_id)496        except ValueError as e:497            if "Unrecognized configuration class" in str(e):498                self.model = AutoModelForImageTextToText.from_pretrained(model_id, device_map=device_map)499                self.processor = AutoProcessor.from_pretrained(model_id)500                self._is_vlm = True501            else:502                raise e503        except Exception as e:504            logger.warning(505                f"Failed to load tokenizer and model for {model_id=}: {e}. Loading default tokenizer and model instead from {default_model_id=}."506            )507            self.model_id = default_model_id508            self.tokenizer = AutoTokenizer.from_pretrained(default_model_id)509            self.model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device_map, torch_dtype=torch_dtype)510 511    def make_stopping_criteria(self, stop_sequences: List[str], tokenizer) -> "StoppingCriteriaList":512        from transformers import StoppingCriteria, StoppingCriteriaList513 514        class StopOnStrings(StoppingCriteria):515            def __init__(self, stop_strings: List[str], tokenizer):516                self.stop_strings = stop_strings517                self.tokenizer = tokenizer518                self.stream = ""519 520            def reset(self):521                self.stream = ""522 523            def __call__(self, input_ids, scores, **kwargs):524                generated = self.tokenizer.decode(input_ids[0][-1], skip_special_tokens=True)525                self.stream += generated526                if any([self.stream.endswith(stop_string) for stop_string in self.stop_strings]):527                    return True528                return False529 530        return StoppingCriteriaList([StopOnStrings(stop_sequences, tokenizer)])531 532    def __call__(533        self,534        messages: List[Dict[str, str]],535        stop_sequences: Optional[List[str]] = None,536        grammar: Optional[str] = None,537        tools_to_call_from: Optional[List[Tool]] = None,538        images: Optional[List[Image.Image]] = None,539        **kwargs,540    ) -> ChatMessage:541        completion_kwargs = self._prepare_completion_kwargs(542            messages=messages,543            stop_sequences=stop_sequences,544            grammar=grammar,545            flatten_messages_as_text=(not self._is_vlm),546            **kwargs,547        )548 549        messages = completion_kwargs.pop("messages")550        stop_sequences = completion_kwargs.pop("stop", None)551 552        max_new_tokens = (553            kwargs.get("max_new_tokens")554            or kwargs.get("max_tokens")555            or self.kwargs.get("max_new_tokens")556            or self.kwargs.get("max_tokens")557        )558 559        if max_new_tokens:560            completion_kwargs["max_new_tokens"] = max_new_tokens561 562        if hasattr(self, "processor"):563            images = [Image.open(image) for image in images] if images else None564            prompt_tensor = self.processor.apply_chat_template(565                messages,566                tools=[get_tool_json_schema(tool) for tool in tools_to_call_from] if tools_to_call_from else None,567                return_tensors="pt",568                tokenize=True,569                return_dict=True,570                images=images,571                add_generation_prompt=True if tools_to_call_from else False,572            )573        else:574            prompt_tensor = self.tokenizer.apply_chat_template(575                messages,576                tools=[get_tool_json_schema(tool) for tool in tools_to_call_from] if tools_to_call_from else None,577                return_tensors="pt",578                return_dict=True,579                add_generation_prompt=True if tools_to_call_from else False,580            )581 582        prompt_tensor = prompt_tensor.to(self.model.device)583        count_prompt_tokens = prompt_tensor["input_ids"].shape[1]584 585        if stop_sequences:586            stopping_criteria = self.make_stopping_criteria(587                stop_sequences, tokenizer=self.processor if hasattr(self, "processor") else self.tokenizer588            )589        else:590            stopping_criteria = None591 592        out = self.model.generate(593            **prompt_tensor,594            stopping_criteria=stopping_criteria,595            **completion_kwargs,596        )597        generated_tokens = out[0, count_prompt_tokens:]598        if hasattr(self, "processor"):599            output = self.processor.decode(generated_tokens, skip_special_tokens=True)600        else:601            output = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)602        self.last_input_token_count = count_prompt_tokens603        self.last_output_token_count = len(generated_tokens)604 605        if stop_sequences is not None:606            output = remove_stop_sequences(output, stop_sequences)607 608        if tools_to_call_from is None:609            return ChatMessage(610                role="assistant",611                content=output,612                raw={"out": out, "completion_kwargs": completion_kwargs},613            )614        else:615            if "Action:" in output:616                output = output.split("Action:", 1)[1].strip()617            try:618                start_index = output.index("{")619                end_index = output.rindex("}")620                output = output[start_index : end_index + 1]621            except Exception as e:622                raise Exception("No json blob found in output!") from e623 624            try:625                parsed_output = json.loads(output)626            except json.JSONDecodeError as e:627                raise ValueError(f"Tool call '{output}' has an invalid JSON structure: {e}")628            tool_name = parsed_output.get("name")629            tool_arguments = parsed_output.get("arguments")630            return ChatMessage(631                role="assistant",632                content="",633                tool_calls=[634                    ChatMessageToolCall(635                        id="".join(random.choices("0123456789", k=5)),636                        type="function",637                        function=ChatMessageToolCallDefinition(name=tool_name, arguments=tool_arguments),638                    )639                ],640                raw={"out": out, "completion_kwargs": completion_kwargs},641            )642 643 644class LiteLLMModel(Model):645    """This model connects to [LiteLLM](https://www.litellm.ai/) as a gateway to hundreds of LLMs.646 647    Parameters:648        model_id (`str`):649            The model identifier to use on the server (e.g. "gpt-3.5-turbo").650        api_base (`str`, *optional*):651            The base URL of the OpenAI-compatible API server.652        api_key (`str`, *optional*):653            The API key to use for authentication.654        custom_role_conversions (`dict[str, str]`, *optional*):655            Custom role conversion mapping to convert message roles in others.656            Useful for specific models that do not support specific message roles like "system".657        **kwargs:658            Additional keyword arguments to pass to the OpenAI API.659    """660 661    def __init__(662        self,663        model_id: str = "anthropic/claude-3-5-sonnet-20240620",664        api_base=None,665        api_key=None,666        custom_role_conversions: Optional[Dict[str, str]] = None,667        **kwargs,668    ):669        try:670            import litellm671        except ModuleNotFoundError:672            raise ModuleNotFoundError(673                "Please install 'litellm' extra to use LiteLLMModel: `pip install 'smolagents[litellm]'`"674            )675 676        super().__init__(**kwargs)677        self.model_id = model_id678        # IMPORTANT - Set this to TRUE to add the function to the prompt for Non OpenAI LLMs679        litellm.add_function_to_prompt = True680        self.api_base = api_base681        self.api_key = api_key682        self.custom_role_conversions = custom_role_conversions683 684    def __call__(685        self,686        messages: List[Dict[str, str]],687        stop_sequences: Optional[List[str]] = None,688        grammar: Optional[str] = None,689        tools_to_call_from: Optional[List[Tool]] = None,690        **kwargs,691    ) -> ChatMessage:692        import litellm693 694        completion_kwargs = self._prepare_completion_kwargs(695            messages=messages,696            stop_sequences=stop_sequences,697            grammar=grammar,698            tools_to_call_from=tools_to_call_from,699            model=self.model_id,700            api_base=self.api_base,701            api_key=self.api_key,702            convert_images_to_image_urls=True,703            flatten_messages_as_text=self.model_id.startswith("ollama"),704            custom_role_conversions=self.custom_role_conversions,705            **kwargs,706        )707 708        response = litellm.completion(**completion_kwargs)709 710        self.last_input_token_count = response.usage.prompt_tokens711        self.last_output_token_count = response.usage.completion_tokens712        message = ChatMessage.from_dict(713            response.choices[0].message.model_dump(include={"role", "content", "tool_calls"})714        )715        message.raw = response716 717        if tools_to_call_from is not None:718            return parse_tool_args_if_needed(message)719        return message720 721 722class OpenAIServerModel(Model):723    """This model connects to an OpenAI-compatible API server.724 725    Parameters:726        model_id (`str`):727            The model identifier to use on the server (e.g. "gpt-3.5-turbo").728        api_base (`str`, *optional*):729            The base URL of the OpenAI-compatible API server.730        api_key (`str`, *optional*):731            The API key to use for authentication.732        organization (`str`, *optional*):733            The organization to use for the API request.734        project (`str`, *optional*):735            The project to use for the API request.736        custom_role_conversions (`dict[str, str]`, *optional*):737            Custom role conversion mapping to convert message roles in others.738            Useful for specific models that do not support specific message roles like "system".739        **kwargs:740            Additional keyword arguments to pass to the OpenAI API.741    """742 743    def __init__(744        self,745        model_id: str,746        api_base: Optional[str] = None,747        api_key: Optional[str] = None,748        organization: Optional[str] | None = None,749        project: Optional[str] | None = None,750        custom_role_conversions: Optional[Dict[str, str]] = None,751        **kwargs,752    ):753        try:754            import openai755        except ModuleNotFoundError:756            raise ModuleNotFoundError(757                "Please install 'openai' extra to use OpenAIServerModel: `pip install 'smolagents[openai]'`"758            ) from None759 760        super().__init__(**kwargs)761        self.model_id = model_id762        self.client = openai.OpenAI(763            base_url=api_base,764            api_key=api_key,765            organization=organization,766            project=project,767        )768        self.custom_role_conversions = custom_role_conversions769 770    def __call__(771        self,772        messages: List[Dict[str, str]],773        stop_sequences: Optional[List[str]] = None,774        grammar: Optional[str] = None,775        tools_to_call_from: Optional[List[Tool]] = None,776        **kwargs,777    ) -> ChatMessage:778        completion_kwargs = self._prepare_completion_kwargs(779            messages=messages,780            stop_sequences=stop_sequences,781            grammar=grammar,782            tools_to_call_from=tools_to_call_from,783            model=self.model_id,784            custom_role_conversions=self.custom_role_conversions,785            convert_images_to_image_urls=True,786            **kwargs,787        )788        response = self.client.chat.completions.create(**completion_kwargs)789        self.last_input_token_count = response.usage.prompt_tokens790        self.last_output_token_count = response.usage.completion_tokens791 792        message = ChatMessage.from_dict(793            response.choices[0].message.model_dump(include={"role", "content", "tool_calls"})794        )795        message.raw = response796        if tools_to_call_from is not None:797            return parse_tool_args_if_needed(message)798        return message799 800 801class AzureOpenAIServerModel(OpenAIServerModel):802    """This model connects to an Azure OpenAI deployment.803 804    Parameters:805        model_id (`str`):806            The model deployment name to use when connecting (e.g. "gpt-4o-mini").807        azure_endpoint (`str`, *optional*):808            The Azure endpoint, including the resource, e.g. `https://example-resource.azure.openai.com/`. If not provided, it will be inferred from the `AZURE_OPENAI_ENDPOINT` environment variable.809        api_key (`str`, *optional*):810            The API key to use for authentication. If not provided, it will be inferred from the `AZURE_OPENAI_API_KEY` environment variable.811        api_version (`str`, *optional*):812            The API version to use. If not provided, it will be inferred from the `OPENAI_API_VERSION` environment variable.813        custom_role_conversions (`dict[str, str]`, *optional*):814            Custom role conversion mapping to convert message roles in others.815            Useful for specific models that do not support specific message roles like "system".816        **kwargs:817            Additional keyword arguments to pass to the Azure OpenAI API.818    """819 820    def __init__(821        self,822        model_id: str,823        azure_endpoint: Optional[str] = None,824        api_key: Optional[str] = None,825        api_version: Optional[str] = None,826        custom_role_conversions: Optional[Dict[str, str]] = None,827        **kwargs,828    ):829        # read the api key manually, to avoid super().__init__() trying to use the wrong api_key (OPENAI_API_KEY)830        if api_key is None:831            api_key = os.environ.get("AZURE_OPENAI_API_KEY")832 833        super().__init__(model_id=model_id, api_key=api_key, custom_role_conversions=custom_role_conversions, **kwargs)834        # if we've reached this point, it means the openai package is available (checked in baseclass) so go ahead and import it835        import openai836 837        self.client = openai.AzureOpenAI(api_key=api_key, api_version=api_version, azure_endpoint=azure_endpoint)838 839 840__all__ = [841    "MessageRole",842    "tool_role_conversions",843    "get_clean_message_list",844    "Model",845    "TransformersModel",846    "HfApiModel",847    "LiteLLMModel",848    "OpenAIServerModel",849    "AzureOpenAIServerModel",850    "ChatMessage",851]