dkolarova/bpm-agent
0
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]