From 39decab96a98fe7a358d34796af22469b2084dbb Mon Sep 17 00:00:00 2001 From: Madeesh Kannan Date: Mon, 19 Aug 2024 15:33:48 +0200 Subject: [PATCH] refactor: Remove usage of deprecated `ChatMessage.to_openai_format` (#1007) --- .../amazon_bedrock/chat/adapters.py | 20 ++++--------------- 1 file changed, 4 insertions(+), 16 deletions(-) diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/adapters.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/adapters.py index 87bd61b53..56eefdf09 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/adapters.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/adapters.py @@ -5,6 +5,7 @@ from typing import Any, Callable, ClassVar, Dict, List, Optional from botocore.eventstream import EventStream +from haystack.components.generators.openai_utils import _convert_message_to_openai_format from haystack.dataclasses import ChatMessage, ChatRole, StreamingChunk from transformers import AutoTokenizer, PreTrainedTokenizer @@ -389,27 +390,14 @@ def prepare_chat_messages(self, messages: List[ChatMessage]) -> str: # default is https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1/blob/main/tokenizer_config.json # but we'll use our custom chat template prepared_prompt: str = self.prompt_handler.tokenizer.apply_chat_template( - conversation=[self.to_openai_format(m) for m in messages], tokenize=False, chat_template=self.chat_template + conversation=[_convert_message_to_openai_format(m) for m in messages], + tokenize=False, + chat_template=self.chat_template, ) if self.truncate: prepared_prompt = self._ensure_token_limit(prepared_prompt) return prepared_prompt - def to_openai_format(self, m: ChatMessage) -> Dict[str, Any]: - """ - Convert the message to the format expected by OpenAI's Chat API. - See the [API reference](https://platform.openai.com/docs/api-reference/chat/create) for details. - - :returns: A dictionary with the following key: - - `role` - - `content` - - `name` (optional) - """ - msg = {"role": m.role.value, "content": m.content} - if m.name: - msg["name"] = m.name - return msg - def check_prompt(self, prompt: str) -> Dict[str, Any]: """ Checks the prompt length and resizes it if necessary. If the prompt is too long, it will be truncated.